Implement finding

Author: henryk

fints/formals.py

@@ -3,7 +3,7 @@
 from contextlib import suppress
 from inspect import getmro
 from copy import deepcopy
-from collections import OrderedDict
+from collections import OrderedDict, Iterable
 
 from fints.utils import classproperty, SubclassesMixin, Password
 
@@ -392,11 +392,42 @@ def find_segments(self, type=None, version=None, callback=None, recurse=True):
         The match results of all given parameters will be AND-combined.
         """
 
+        if type is None:
+            type = []
+        elif isinstance(type, str) or not isinstance(type, (list, tuple, Iterable)):
+            type = [type]
+
+        if version is None:
+            version = []
+        elif not isinstance(version, (list, tuple, Iterable)):
+            version = [version]
+
+        if callback is None:
+            callback = lambda s: True
+
+        for s in self.segments:
+            if ((not type) or any(s.header.type == t for t in type)) and \
+                ((not version) or any(s.header.version == v for v in version)) and \
+                callback(s):
+                yield s
+
+            if recurse:
+                for name, field in s._fields.items():
+                    if isinstance(field, SegmentSequenceField):
+                        val = getattr(s, name)
+                        if val:
+                            yield from val.find_segments(type=type, version=version, callback=callback, recurse=recurse)
+
     def find_segment_first(self, *args, **kwargs):
         """Finds the first matching segment.
 
         Same parameters as find_segments(), but only returns the first match, or None if no match is found."""
 
+        for m in self.find_segments(*args, **kwargs):
+            return m
+
+        return None
+
 class SegmentSequenceField(DataElementField):
     type = 'sf'
 

tests/test_formals.py

@@ -333,3 +333,26 @@ class A(Container):
 
     with pytest.warns(UserWarning, match=r'Generic field used for type None value \[\]'):
         i2 = A(a=[])
+
+def test_find_1():
+    from conftest import COMPLICATED_EXAMPLE
+    from fints.parser import FinTS3Parser
+    from fints.segments import HNHBS1, HNHBK3
+
+    m = FinTS3Parser().parse_message(COMPLICATED_EXAMPLE)
+
+    assert len(list(m.find_segments('HNHBK'))) == 1
+    assert len(list(m.find_segments( ['HNHBK', 'HNHBS'] ))) == 2
+
+    assert len(list(m.find_segments( ['HNHBK', 'HNHBS'], 1))) == 1
+    assert len(list(m.find_segments( ['HNHBK', 'HNHBS'], (1, 3) ))) == 2
+
+    assert isinstance(m.find_segment_first('HNHBK'), HNHBK3)
+    assert isinstance(m.find_segment_first('HNHBS'), HNHBS1)
+
+    assert m.find_segment_first('ITST') is None
+
+    assert len( m.find_segment_first('HITANS', 1).parameters.twostep_parameters ) == 2
+    assert len( m.find_segment_first('HITANS', 3).parameters.twostep_parameters ) == 6
+
+    assert m.find_segment_first('HITANS', recurse=False) is None