From 207118629c1f92af22db839a00f91e9196608adf Mon Sep 17 00:00:00 2001
From: cogu <congus8@gmail.com>
Date: Wed, 13 Mar 2024 23:24:51 +0100
Subject: [PATCH] Add new package filter functions

---
 src/autosar/xml/element.py   | 30 +++++++++++++++++++++++++++---
 src/autosar/xml/workspace.py | 22 ++++++++++++++++------
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/src/autosar/xml/element.py b/src/autosar/xml/element.py
index 483815a..d80e661 100644
--- a/src/autosar/xml/element.py
+++ b/src/autosar/xml/element.py
@@ -3258,11 +3258,11 @@ class Package(CollectableElement):
 
     def __init__(self, name: str, **kwargs: dict) -> None:
         super().__init__(name, **kwargs)
-        self.elements = []
-        self.packages = []
+        self.elements: list[ARElement] = []
+        self.packages: list[Package] = []
         self._collection_map = {}
 
-    def append(self, item: Any):
+    def append(self, item: CollectableElement):
         """
         Append element or sub-package
         """
@@ -3327,6 +3327,30 @@ def find(self, ref: str) -> Any:
                 return item.find(parts[2])
         return item
 
+    def filter(self, *names: str) -> Iterator[ARElement]:
+        """
+        Yields all elements whose short-name matches any of the names in
+        argument list
+        """
+        for elem in self.elements:
+            if elem.name in names:
+                yield elem
+
+    def filter_regex(self, pattern: str | re.Pattern) -> Iterator[ARElement]:
+        """
+        Yields all elements whose short-name matches the given regex pattern
+        """
+        regex: re.Pattern = None
+        if isinstance(pattern, str):
+            regex = re.compile(pattern)
+        elif isinstance(pattern, re.Pattern):
+            regex = pattern
+        else:
+            raise TypeError(f"pattern: Invalid type '{str(type(pattern))}'")
+        for elem in self.elements:
+            if regex.match(elem.name):
+                yield elem
+
 # --- ModeDeclaration elements
 
 
diff --git a/src/autosar/xml/workspace.py b/src/autosar/xml/workspace.py
index e8df95a..4062b4d 100644
--- a/src/autosar/xml/workspace.py
+++ b/src/autosar/xml/workspace.py
@@ -130,10 +130,10 @@ def init_package_map(self, mapping: dict[str, str]) -> None:
         """
         Initializes an internally stored package map using key-value pairs.
         The key can be any (unique) name while each value must be a package reference.
-        This function internally creates packages from dict-values using the methd
+        This function internally creates packages from dict-values calling the method
         make_packages.
 
-        Use in conjunction with the add_element function.
+        Use in conjunction with the add_element and find_element methods.
 
         Avoid manually calling make_packages if using this method.
         """
@@ -145,22 +145,32 @@ def add_element(self, package_key: str, element: ar_element.ARObject):
         """
         Adds element to package specified by package_key
 
-        The method init_package_map must be called before using this method.
+        Only use after calling init_package_map.
         """
         if len(self.package_map) == 0:
             raise RuntimeError("Internal package map not initialized")
         self.package_map[package_key].append(element)
 
-    def find_element(self, package_key: str, element_name: str) -> ar_element.Identifiable | None:
+    def find_element(self, package_key: str, element_name: str) -> ar_element.ARElement | None:
         """
-        Finds element in package referenced by package_key.
+        Finds an element in the package referenced by package_key.
 
-        Ony use together with init_package_map and add_element.
+        Only use after calling init_package_map.
         """
         if len(self.package_map) == 0:
             raise RuntimeError("Internal package map not initialized")
         return self.package_map[package_key].find(element_name)
 
+    def get_package(self, package_key: str) -> ar_element.Package:
+        """
+        Returns the package referenced by package_key.
+
+        Only use after calling init_package_map.
+        """
+        if len(self.package_map) == 0:
+            raise RuntimeError("Internal package map not initialized")
+        return self.package_map[package_key]
+
     def append(self, package: ar_element.Package):
         """
         Appends package to this worksapace