From 5047490c4b2cd6e006a4b3d3e866584d81bf510d Mon Sep 17 00:00:00 2001
From: Jocelyn Delalande <jocelyn@crapouillou.net>
Date: Thu, 10 Sep 2020 11:55:01 +0200
Subject: [PATCH] Support extra elements for header segment of message and
 interchange

Fix #19
---
 pydifact/segmentcollection.py   | 18 ++++++++++++++----
 tests/test_segmentcollection.py | 16 ++++++++++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/pydifact/segmentcollection.py b/pydifact/segmentcollection.py
index 60d225c..aff72c1 100644
--- a/pydifact/segmentcollection.py
+++ b/pydifact/segmentcollection.py
@@ -21,7 +21,7 @@
 # THE SOFTWARE.
 
 import collections
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 import datetime
 
 from pydifact.parser import Parser
@@ -34,12 +34,18 @@
 class AbstractSegmentsContainer:
     """Represent a collection of EDI Segments for both reading and writing."""
 
-    def __init__(self):
+    def __init__(self, extra_header_elements: List[Union[str, List[str]]] = []):
+        """
+        :param extra_header_elements: a list of elements to be appended at the end
+          of the header segment (same format as Segment() constructor *elements).
+        """
 
         # The segments that make up this message
         self.segments = []
         self.characters = Characters()
 
+        self.extra_header_elements = extra_header_elements
+
         # Flag whether the UNA header is present
         self.has_una_segment = False
 
@@ -189,8 +195,9 @@ def __init__(
             self,
             reference_number: str,
             identifier: Tuple,
+            *args, **kwargs
     ):
-        super().__init__()
+        super().__init__(*args, **kwargs)
         self.reference_number = reference_number
         self.identifier = identifier
 
@@ -199,6 +206,7 @@ def get_header_segment(self) -> Segment:
             "UNH",
             self.reference_number,
             [str(i) for i in self.identifier],
+            *self.extra_header_elements,
         )
 
     def get_footer_segment(self) -> Segment:
@@ -232,8 +240,9 @@ def __init__(
             syntax_identifier: Tuple[str, int],
             delimiters: Characters = Characters(),
             timestamp: datetime.datetime = None,
+            *args, **kwargs
     ):
-        super().__init__()
+        super().__init__(*args, **kwargs)
         self.sender = sender
         self.recipient = recipient
         self.control_reference = control_reference
@@ -249,6 +258,7 @@ def get_header_segment(self) -> Segment:
             self.recipient,
             ['{:%y%m%d}'.format(self.timestamp), '{:%H%M}'.format(self.timestamp)],
             self.control_reference,
+            *self.extra_header_elements,
         )
 
     def get_footer_segment(self) -> Segment:
diff --git a/tests/test_segmentcollection.py b/tests/test_segmentcollection.py
index 9445967..d2aecc4 100644
--- a/tests/test_segmentcollection.py
+++ b/tests/test_segmentcollection.py
@@ -125,6 +125,22 @@ def test_empty_interchange(interchange):
     )
 
 
+def test_empty_interchange_w_extra_header(interchange):
+    i = Interchange(
+        sender='1234',
+        recipient='3333',
+        timestamp=datetime.datetime(2020, 1, 2, 22, 12),
+        control_reference='42',
+        syntax_identifier=('UNOC', 1),
+        extra_header_elements=[['66', '2'], 'ZZ']
+    )
+
+    assert str(i) == (
+        "UNB+UNOC:1+1234+3333+200102:2212+42+66:2+ZZ'"
+        "UNZ+0+42'"
+    )
+
+
 def test_empty_interchange_from_str():
     i = Interchange.from_str(
         "UNB+UNOC:1+1234+3333+200102:2212+42'"