use Filters instead of kwargs

Author: vignesharivazhagan

py2puml/export/__init__.py

@@ -1,4 +1,5 @@
-from typing import Iterable, List
+from dataclasses import dataclass
+from typing import Callable, Iterable, List, Optional
 
 from py2puml.domain.umlclass import UmlClass
 from py2puml.domain.umlenum import UmlEnum
@@ -26,14 +27,39 @@
 FEATURE_INSTANCE = ''
 
 
+@dataclass
+class Filters:
+    skip_block: Optional[Callable[[UmlItem], bool]] = None
+    skip_relation: Optional[Callable[[UmlRelation], bool]] = None
+
+
+def should_skip(filter: Callable | None, item: UmlItem | UmlRelation) -> bool:
+    if filter is None:
+        return False
+
+    if not callable(filter):
+        raise ValueError('Filter must be a callable')
+
+    try:
+        _should_skip = filter(item)
+        if not isinstance(_should_skip, bool):
+            raise ValueError('Filter must return a boolean value')
+        return _should_skip
+    except Exception as e:
+        raise ValueError('Error while applying filter') from e
+
+
 def to_puml_content(
-    diagram_name: str, uml_items: List[UmlItem], uml_relations: List[UmlRelation], **kwargs
+    diagram_name: str, uml_items: List[UmlItem], uml_relations: List[UmlRelation], filters: Optional[Filters] = None
 ) -> Iterable[str]:
+    if filters is None:
+        filters = Filters()
+
     yield PUML_FILE_START.format(diagram_name=diagram_name)
 
     # exports the domain classes and enums
     for uml_item in uml_items:
-        if 'is_block_valid' in kwargs and callable(kwargs['is_block_valid']) and not kwargs['is_block_valid'](uml_item):
+        if should_skip(filters.skip_block, uml_item):
             continue
         if isinstance(uml_item, UmlEnum):
             uml_enum: UmlEnum = uml_item
@@ -52,26 +78,14 @@ def to_puml_content(
                     attr_type=uml_attr.type,
                     staticity=FEATURE_STATIC if uml_attr.static else FEATURE_INSTANCE,
                 )
-            for uml_method in uml_class.methods:
-                if (
-                    'is_method_valid' in kwargs
-                    and callable(kwargs['is_method_valid'])
-                    and not kwargs['is_method_valid'](uml_method)
-                ):
-                    continue
-
-                yield f'  {uml_method.represent_as_puml()}\n'
+            # TODO: Add skip_method filter here once PR #43 is merged
             yield PUML_ITEM_END
         else:
             raise TypeError(f'cannot process uml_item of type {uml_item.__class__}')
 
     # exports the domain relationships between classes and enums
     for uml_relation in uml_relations:
-        if (
-            'is_relation_valid' in kwargs
-            and callable(kwargs['is_relation_valid'])
-            and not kwargs['is_relation_valid'](uml_relation)
-        ):
+        if should_skip(filters.skip_relation, uml_relation):
             continue
         yield PUML_RELATION_TPL.format(
             source_fqn=uml_relation.source_fqn, rel_type=uml_relation.type.value, target_fqn=uml_relation.target_fqn

py2puml/export/puml.py

@@ -1,14 +1,14 @@
-from typing import Dict, Iterable, List
+from typing import Dict, Iterable, List, Optional
 
 from py2puml.domain.umlitem import UmlItem
 from py2puml.domain.umlrelation import UmlRelation
-from py2puml.export.puml import to_puml_content
+from py2puml.export.puml import Filters, to_puml_content
 from py2puml.inspection.inspectpackage import inspect_package
 
 
-def py2puml(domain_path: str, domain_module: str, **kwargs) -> Iterable[str]:
+def py2puml(domain_path: str, domain_module: str, filters: Optional[Filters] = None) -> Iterable[str]:
     domain_items_by_fqn: Dict[str, UmlItem] = {}
     domain_relations: List[UmlRelation] = []
     inspect_package(domain_path, domain_module, domain_items_by_fqn, domain_relations)
 
-    return to_puml_content(domain_module, domain_items_by_fqn.values(), domain_relations, **kwargs)
+    return to_puml_content(domain_module, domain_items_by_fqn.values(), domain_relations, filters)