1919# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
2020# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
2121# THE SOFTWARE.
22- from typing import Optional , Generator , Any
22+ from collections . abc import Iterator
2323
2424from pydifact .tokenizer import Tokenizer
2525from pydifact .token import Token
26- from pydifact .segments import Segment , SegmentFactory
26+ from pydifact .segments import Element , Elements , Segment , SegmentFactory
2727from pydifact .control import Characters
2828
2929
@@ -32,15 +32,15 @@ class Parser:
3232
3333 def __init__ (
3434 self ,
35- factory : Optional [ SegmentFactory ] = None ,
36- characters : Optional [ Characters ] = None ,
37- ):
35+ factory : SegmentFactory | None = None ,
36+ characters : Characters | None = None ,
37+ ) -> None :
3838 self .factory = factory or SegmentFactory ()
3939 self .characters = characters or Characters ()
4040
4141 def parse (
42- self , message : str , characters : Characters = None
43- ) -> Generator [Segment , Any , None ]:
42+ self , message : str , characters : Characters | None = None
43+ ) -> Iterator [Segment ]:
4444 """Parse the message into a list of segments.
4545
4646 :param characters: the control characters to use, if there is no
@@ -87,8 +87,8 @@ def parse(
8787
8888 @staticmethod
8989 def get_control_characters (
90- message : str , characters : Characters = None
91- ) -> Optional [ Characters ] :
90+ message : str , characters : Characters | None = None
91+ ) -> Characters | None :
9292 """Read the UNA segment from the passed string and extract/store the control characters from it.
9393
9494 :param message: a valid EDI message string, or UNA segment string,
@@ -112,7 +112,6 @@ def get_control_characters(
112112
113113 # Get the character definitions
114114 chars = message [3 :9 ]
115- characters .is_extracted_from_message = True
116115
117116 characters .component_separator = chars [0 ]
118117 characters .data_separator = chars [1 ]
@@ -124,8 +123,8 @@ def get_control_characters(
124123 return characters
125124
126125 def convert_tokens_to_segments (
127- self , tokens : list , characters : Characters , with_una : bool = False
128- ):
126+ self , tokens : Iterator [ Token ] , characters : Characters , with_una : bool = False
127+ ) -> Iterator [ Segment ] :
129128 """Convert the tokenized message into an array of segments.
130129 :param tokens: The tokens that make up the message
131130 :param characters: the control characters to use
@@ -134,9 +133,10 @@ def convert_tokens_to_segments(
134133 :rtype list of Segment
135134 """
136135
137- segments = []
138- current_segment = []
139- data_element = None
136+ segments : list [Elements ] = []
137+ current_segment : Elements = []
138+ data_element : list [str ] = []
139+ data_element_value : Element
140140 in_segment = False
141141 empty_component_counter = 0
142142
@@ -149,12 +149,14 @@ def convert_tokens_to_segments(
149149 if token .type == Token .Type .TERMINATOR :
150150 in_segment = False
151151 if len (data_element ) == 0 : # empty element
152- data_element = ""
153- if len (data_element ) == 1 :
152+ data_element_value = ""
153+ elif len (data_element ) == 1 :
154154 # use a str instead of a list
155- data_element = data_element [0 ]
155+ data_element_value = data_element [0 ]
156+ else :
157+ data_element_value = data_element
156158
157- current_segment .append (data_element )
159+ current_segment .append (data_element_value )
158160 data_element = []
159161 continue
160162
@@ -174,11 +176,13 @@ def convert_tokens_to_segments(
174176 # data_element to an empty list []
175177 if token .type == Token .Type .DATA_SEPARATOR :
176178 if len (data_element ) == 0 : # empty element
177- data_element = ""
179+ data_element_value = ""
178180 elif len (data_element ) == 1 :
179- data_element = data_element [0 ]
181+ data_element_value = data_element [0 ]
182+ else :
183+ data_element_value = data_element
180184
181- current_segment .append (data_element )
185+ current_segment .append (data_element_value )
182186
183187 data_element = []
184188 empty_component_counter = 0
@@ -210,4 +214,5 @@ def convert_tokens_to_segments(
210214
211215 for segment in segments :
212216 name = segment .pop (0 )
213- yield self .factory .create_segment (name , * segment )
217+ # create_segment tests name type
218+ yield self .factory .create_segment (name , * segment ) # type: ignore
0 commit comments