Add write option to override missing transfer syntax in dataset

read_test.go

@@ -606,7 +606,10 @@ type headerData struct {
 // Write a collection of elements and return them as an encoded buffer of bytes.
 func writeElements(elements []*Element) ([]byte, error) {
 	buff := bytes.Buffer{}
-	dcmWriter := NewWriter(&buff)
+	dcmWriter, err := NewWriter(&buff)
+	if err != nil {
+		return []byte{}, err
+	}
 	dcmWriter.SetTransferSyntax(binary.LittleEndian, true)
 
 	for _, e := range elements {

write.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"slices"
 
 	"github.com/suyashkumar/dicom/pkg/vrraw"
@@ -37,14 +38,17 @@ type Writer struct {
 }
 
 // NewWriter returns a new Writer, that points to the provided io.Writer.
-func NewWriter(out io.Writer, opts ...WriteOption) *Writer {
+func NewWriter(out io.Writer, opts ...WriteOption) (*Writer, error) {
 	optSet := toWriteOptSet(opts...)
+	if err := optSet.validate(); err != nil {
+		return nil, err
+	}
 	w := dicomio.NewWriter(out, nil, false)
 
 	return &Writer{
 		writer: w,
 		optSet: optSet,
-	}
+	}, nil
 }
 
 // SetTransferSyntax sets the transfer syntax for the underlying dicomio.Writer.
@@ -66,17 +70,23 @@ func (w *Writer) writeDataset(ds Dataset) error {
 		return err
 	}
 
-	endian, implicit, err := ds.transferSyntax()
-	if (err != nil && err != ErrorElementNotFound) || (err == ErrorElementNotFound && !w.optSet.defaultMissingTransferSyntax) {
+	bo, implicit, err := ds.transferSyntax()
+	if (err != nil && err != ErrorElementNotFound) || (err == ErrorElementNotFound && (!w.optSet.defaultMissingTransferSyntax && w.optSet.overrideMissingTransferSyntaxUID == "")) {
 		return err
 	}
 
-	if err == ErrorElementNotFound && w.optSet.defaultMissingTransferSyntax {
-		w.writer.SetTransferSyntax(binary.LittleEndian, true)
-	} else {
-		w.writer.SetTransferSyntax(endian, implicit)
+	if errors.Is(err, ErrorElementNotFound) && w.optSet.defaultMissingTransferSyntax {
+		bo = binary.LittleEndian
+		implicit = true
+	} else if errors.Is(err, ErrorElementNotFound) && w.optSet.overrideMissingTransferSyntaxUID != "" {
+		bo, implicit, err = uid.ParseTransferSyntaxUID(w.optSet.overrideMissingTransferSyntaxUID)
+		if err != nil {
+			return err
+		}
 	}
 
+	w.writer.SetTransferSyntax(bo, implicit)
+
 	for _, elem := range ds.Elements {
 		if elem.Tag.Group != tag.MetadataGroup {
 			err = writeElement(w.writer, elem, *w.optSet)
@@ -97,7 +107,10 @@ func (w *Writer) WriteElement(e *Element) error {
 // Write will write the input DICOM dataset to the provided io.Writer as a complete DICOM (including any header
 // information if available).
 func Write(out io.Writer, ds Dataset, opts ...WriteOption) error {
-	w := NewWriter(out, opts...)
+	w, err := NewWriter(out, opts...)
+	if err != nil {
+		return err
+	}
 	return w.writeDataset(ds)
 }
 
@@ -130,11 +143,33 @@ func DefaultMissingTransferSyntax() WriteOption {
 	}
 }
 
+// OverrideMissingTransferSyntaxWith returns a WriteOption indicating that if
+// the dicom to be written does _not_ have a transfer syntax UID in its metadata
+// that it should be written using the provided transferSyntaxUID.
+//
+// If the Writer is unable to recognize or write out using the provided
+// transferSyntaxUID, an error will be returned at initialization time.
+func OverrideMissingTransferSyntaxWith(transferSyntaxUID string) WriteOption {
+	return func(set *writeOptSet) {
+		set.overrideMissingTransferSyntaxUID = transferSyntaxUID
+	}
+}
+
 // writeOptSet represents the flattened option set after all WriteOptions have been applied.
 type writeOptSet struct {
-	skipVRVerification           bool
-	skipValueTypeVerification    bool
-	defaultMissingTransferSyntax bool
+	skipVRVerification               bool
+	skipValueTypeVerification        bool
+	defaultMissingTransferSyntax     bool
+	overrideMissingTransferSyntaxUID string
+}
+
+func (w *writeOptSet) validate() error {
+	if w.overrideMissingTransferSyntaxUID != "" {
+		if _, _, err := uid.ParseTransferSyntaxUID(w.overrideMissingTransferSyntaxUID); err != nil {
+			return fmt.Errorf("unable to parse OverrideMissingTransferSyntaxWith transfer syntax uid %v due to: %s", w.overrideMissingTransferSyntaxUID, err)
+		}
+	}
+	return nil
 }
 
 func toWriteOptSet(opts ...WriteOption) *writeOptSet {
@@ -166,17 +201,31 @@ func writeFileHeader(w *dicomio.Writer, ds *Dataset, metaElems []*Element, opts
 	if err != nil && err != ErrorElementNotFound {
 		return err
 	}
+	log.Println("TEST")
+	log.Println(opts.overrideMissingTransferSyntaxUID)
+
 	err = writeMetaElem(subWriter, tag.TransferSyntaxUID, ds, &tagsUsed, opts)
-	if err != nil && err != ErrorElementNotFound || err == ErrorElementNotFound && !opts.defaultMissingTransferSyntax {
+
+	// TODO(suyashkumar): clean up this logic
+	if err != nil && err != ErrorElementNotFound || (err == ErrorElementNotFound && (!opts.defaultMissingTransferSyntax && opts.overrideMissingTransferSyntaxUID == "")) {
 		return err
 	}
+
 	if err == ErrorElementNotFound && opts.defaultMissingTransferSyntax {
 		// Write the default transfer syntax
 		if err = writeElement(subWriter, mustNewElement(tag.TransferSyntaxUID, []string{uid.ImplicitVRLittleEndian}), opts); err != nil {
 			return err
 		}
 	}
 
+	if err == ErrorElementNotFound && opts.overrideMissingTransferSyntaxUID != "" {
+		// Write the override transfer syntax
+		if err = writeElement(subWriter, mustNewElement(tag.TransferSyntaxUID, []string{opts.overrideMissingTransferSyntaxUID}), opts); err != nil {
+			return err
+		}
+		log.Println("OVERRIDDEN")
+	}
+
 	for _, elem := range metaElems {
 		if elem.Tag.Group == tag.MetadataGroup {
 			if _, ok := tagsUsed[elem.Tag]; !ok {

write_test.go

@@ -866,7 +866,10 @@ func TestWriteElement(t *testing.T) {
 	}}
 
 	buf := bytes.Buffer{}
-	w := NewWriter(&buf)
+	w, err := NewWriter(&buf)
+	if err != nil {
+		t.Fatalf("NewWriter() returned unexpected error: %v", err)
+	}
 	w.SetTransferSyntax(binary.LittleEndian, true)
 
 	for _, e := range writeDS.Elements {
@@ -892,3 +895,65 @@ func TestWriteElement(t *testing.T) {
 		}
 	}
 }
+
+func TestWrite_OverrideMissingTransferSyntaxWith(t *testing.T) {
+	ds := Dataset{Elements: []*Element{
+		mustNewElement(tag.MediaStorageSOPClassUID, []string{"1.2.840.10008.5.1.4.1.1.1.2"}),
+		mustNewElement(tag.MediaStorageSOPInstanceUID, []string{"1.2.3.4.5.6.7"}),
+		mustNewElement(tag.PatientName, []string{"Bob", "Jones"}),
+		mustNewElement(tag.Rows, []int{128}),
+		mustNewElement(tag.FloatingPointValue, []float64{128.10}),
+		mustNewElement(tag.DimensionIndexPointer, []int{32, 36950}),
+		mustNewElement(tag.RedPaletteColorLookupTableData, []byte{0x1, 0x2, 0x3, 0x4}),
+	}}
+
+	cases := []struct {
+		name                   string
+		overrideTransferSyntax string
+	}{
+		{
+			name:                   "Little Endian Implicit",
+			overrideTransferSyntax: "1.2.840.10008.1.2",
+		},
+		{
+			name:                   "Little Endian Explicit",
+			overrideTransferSyntax: "1.2.840.10008.1.2.1",
+		},
+		{
+			name:                   "Big Endian Explicit",
+			overrideTransferSyntax: "1.2.840.10008.1.2.2",
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			writtenDICOM := &bytes.Buffer{}
+			if err := Write(writtenDICOM, ds, OverrideMissingTransferSyntaxWith(tc.overrideTransferSyntax)); err != nil {
+				t.Errorf("Write(OverrideMissingTransferSyntaxWith(%v)) returned unexpected error: %v", tc.overrideTransferSyntax, err)
+			}
+
+			// Read dataset back in to see if no roundtrip errors, and also
+			// check that the written out transfer syntax tag matches.
+
+			parsedDS, err := ParseUntilEOF(writtenDICOM, nil)
+			if err != nil {
+				t.Fatalf("ParseUntilEOF returned unexpected error when reading written dataset back in: %v", err)
+			}
+			tsElem, err := parsedDS.FindElementByTag(tag.TransferSyntaxUID)
+			if err != nil {
+				t.Fatalf("unable to find transfer syntax uid element in written dataset: %v", err)
+			}
+			tsVal, ok := tsElem.Value.GetValue().([]string)
+			if !ok {
+				t.Fatalf("TransferSyntaxUID tag was not of type []")
+			}
+			if len(tsVal) != 1 {
+				t.Errorf("TransferSyntaxUID []string contained more than one element unexpectedly: %v", tsVal)
+			}
+			if tsVal[0] != tc.overrideTransferSyntax {
+				t.Errorf("TransferSyntaxUID in written dicom did not contain the override transfer syntax value. got: %v, want: %v", tsVal[0], tc.overrideTransferSyntax)
+			}
+
+		})
+	}
+}