Extracted traversal rules

Rules/actions.py

@@ -10,6 +10,7 @@
 # Our main goal is to define actions that can be taken on parameters.
 # We'll start by creating a base class that all specific actions will inherit from.
 
+
 class ParameterAction(ABC):
     """
     Base class for actions on parameters.
@@ -40,6 +41,7 @@ def report(self, automate_context: AutomationContext) -> None:
 # Now, let's create a specific action. In this example, we're removing parameters
 # that start with a specific prefix.
 
+
 class PrefixRemovalAction(ParameterAction):
     """
     Action to remove parameters with a specific prefix.
@@ -50,7 +52,9 @@ def __init__(self, forbidden_prefix: str) -> None:
         # The prefix we want to remove
         self.forbidden_prefix: str = forbidden_prefix
 
-    def apply(self, parameter: Union[Base, Dict[str, Any]], parent_object: Base) -> None:
+    def apply(
+        self, parameter: Union[Base, Dict[str, Any]], parent_object: Base
+    ) -> None:
         """
         Remove the parameter if its name starts with the forbidden prefix.
         This function demonstrates the complexities of navigating and modifying
@@ -74,16 +78,16 @@ def apply(self, parameter: Union[Base, Dict[str, Any]], parent_object: Base) ->
 
         # If the parameter's parent object is a Base type, we can use the `__dict__` method
         # to access its underlying dictionary representation.
-        if isinstance(parent_object['parameters'], Base):
+        if isinstance(parent_object["parameters"], Base):
             try:
                 # Retrieve the unique GUID which corresponds to the parameter's key in the parent object.
                 application_name = parameter.__getitem__("applicationInternalName")
 
                 # Check if the GUID exists as a key in the parent object's parameters.
                 # If it does, remove that parameter from the dictionary.
-                if application_name in parent_object['parameters'].__dict__:
-                    parent_object['parameters'].__dict__.pop(application_name)
-                    self.affected_parameters[parent_object['id']].append(param_name)
+                if application_name in parent_object["parameters"].__dict__:
+                    parent_object["parameters"].__dict__.pop(application_name)
+                    self.affected_parameters[parent_object["id"]].append(param_name)
 
             except KeyError:
                 pass
@@ -104,8 +108,7 @@ def report(self, automate_context: AutomationContext) -> None:
 
         # Summarize the names of all removed parameters
         removed_params = set(
-            param for params in self.affected_parameters.values()
-            for param in params
+            param for params in self.affected_parameters.values() for param in params
         )
 
         message = f"The following parameters were removed: {', '.join(removed_params)}"
@@ -114,7 +117,7 @@ def report(self, automate_context: AutomationContext) -> None:
         automate_context.attach_info_to_objects(
             category="Removed_Parameters",
             object_ids=list(self.affected_parameters.keys()),
-            message=message
+            message=message,
         )
 
 
@@ -137,13 +140,12 @@ def apply(self, parameter: Dict[str, str], parent_object: Dict[str, str]) -> Non
         if ParameterRules.has_missing_value(parameter):
             # If missing, add the parameter's name to our affected_parameters dictionary
             # The key is the parent object's ID for easy lookup later
-            self.affected_parameters[parent_object['id']].append(parameter["name"])
+            self.affected_parameters[parent_object["id"]].append(parameter["name"])
 
     def report(self, automate_context: AutomationContext) -> None:
         # Construct a set of unique parameter names that have missing values
         missing_value_params = set(
-            param for params in self.affected_parameters.values()
-            for param in params
+            param for params in self.affected_parameters.values() for param in params
         )
 
         # Formulate a message summarizing the missing parameters
@@ -153,7 +155,7 @@ def report(self, automate_context: AutomationContext) -> None:
         automate_context.attach_info_to_objects(
             category="Missing_Values",
             object_ids=list(self.affected_parameters.keys()),
-            message=message
+            message=message,
         )
 
 
@@ -174,16 +176,15 @@ def apply(self, parameter: Dict[str, str], parent_object: Dict[str, str]) -> Non
         # Check if the parameter has a default value
         if ParameterRules.has_default_value(parameter):
             # If it does, update its value with the new specified value
-            parameter['value'] = self.new_value  # Mutate the parameter value
+            parameter["value"] = self.new_value  # Mutate the parameter value
 
             # Record the parameter's name in our affected_parameters dictionary
-            self.affected_parameters[parent_object['id']].append(parameter["name"])
+            self.affected_parameters[parent_object["id"]].append(parameter["name"])
 
     def report(self, automate_context: AutomationContext) -> None:
         # Construct a set of unique parameter names that were mutated
         mutated_params = set(
-            param for params in self.affected_parameters.values()
-            for param in params
+            param for params in self.affected_parameters.values() for param in params
         )
 
         # Formulate a message summarizing the mutated parameters
@@ -193,5 +194,5 @@ def report(self, automate_context: AutomationContext) -> None:
         automate_context.attach_info_to_objects(
             category="Updated_Defaults",
             object_ids=list(self.affected_parameters.keys()),
-            message=message
+            message=message,
         )

Rules/rules.py

@@ -8,6 +8,7 @@
 # process parameters in our Speckle objects. These rules will be encapsulated
 # in a class called `ParameterRules`.
 
+
 class ParameterRules:
     """
     A collection of rules for processing parameters in Speckle objects.
@@ -23,7 +24,9 @@ def speckle_type_rule(desired_type: str) -> Callable[[Base], bool]:
         """
         Rule: Check if a parameter's speckle_type matches the desired type.
         """
-        return lambda parameter: getattr(parameter, "speckle_type", None) == desired_type
+        return (
+            lambda parameter: getattr(parameter, "speckle_type", None) == desired_type
+        )
 
     @staticmethod
     def forbidden_prefix_rule(given_prefix: str) -> Callable[[Base], bool]:
@@ -50,7 +53,7 @@ def has_missing_value(parameter: Dict[str, str]) -> bool:
         This rule checks if a parameter is missing its value, potentially indicating
         an oversight during data entry or transfer.
         """
-        return not parameter.get('value')
+        return not parameter.get("value")
 
     @staticmethod
     def has_default_value(parameter: Dict[str, str]) -> bool:
@@ -61,7 +64,7 @@ def has_default_value(parameter: Dict[str, str]) -> bool:
         This rule identifies parameters that still have their default values, helping
         to highlight areas where real, meaningful values need to be provided.
         """
-        return parameter.get('value') == "Default"
+        return parameter.get("value") == "Default"
 
     @staticmethod
     def parameter_exists(parameter_name: str, parent_object: Dict[str, str]) -> bool:
@@ -72,4 +75,4 @@ def parameter_exists(parameter_name: str, parent_object: Dict[str, str]) -> bool
         This rule verifies if a specific parameter exists within an object, allowing
         teams to ensure that key data points are always present.
         """
-        return parameter_name in parent_object.get('parameters', {})
+        return parameter_name in parent_object.get("parameters", {})

Rules/traversal.py

@@ -0,0 +1,49 @@
+from specklepy.objects.graph_traversal.traversal import GraphTraversal, TraversalRule
+
+
+def get_data_traversal_rules() -> GraphTraversal:
+    """
+    Generates traversal rules for navigating Speckle data structures.
+
+    This function defines and returns traversal rules tailored for Speckle data.
+    These rules are used to navigate and extract specific data properties
+    within complex Speckle data hierarchies.
+
+    It defines two main rules:
+
+    1. `display_value_rule`:
+        - Targets objects that have properties named either "displayValue" or
+          "@displayValue".
+        - Specifically focuses on objects with a 'speckle_type' containing
+          "Geometry".
+        - For such objects, the function looks to traverse their 'elements'
+          or '@elements' properties.
+
+    2. `default_rule`:
+        - A more general rule that applies to all objects.
+        - It aims to traverse all member names of an object while avoiding
+          deprecated members (a potential enhancement for the future).
+
+    Returns:
+        GraphTraversal: A GraphTraversal instance initialized with the
+        defined rules.
+    """
+    display_value_property_aliases = {"displayValue", "@displayValue"}
+    elements_property_aliases = {"elements", "@elements"}
+
+    display_value_rule = TraversalRule(
+        [
+            lambda o: any(
+                getattr(o, alias, None) for alias in display_value_property_aliases
+            ),
+            lambda o: "Geometry" in o.speckle_type,
+        ],
+        lambda o: elements_property_aliases,
+    )
+
+    default_rule = TraversalRule(
+        [lambda _: True],
+        lambda o: o.get_member_names(),  # TODO: avoid deprecated members
+    )
+
+    return GraphTraversal([display_value_rule, default_rule])

main.py

@@ -11,10 +11,10 @@
     execute_automate_function,
 )
 from specklepy.objects import Base
-from specklepy.objects.graph_traversal.traversal import TraversalRule, GraphTraversal
 
-from Rules.actions import ParameterAction, PrefixRemovalAction, MissingValueReportAction
+from Rules.actions import PrefixRemovalAction
 from Rules.rules import ParameterRules
+from Rules.traversal import get_data_traversal_rules
 
 
 class FunctionInputs(AutomateBase):
@@ -54,12 +54,16 @@ def automate_function(
     version_root_object = automate_context.receive_version()
 
     # Traverse the received Speckle data.
-    speckle_data = get_data_traversal()
+    speckle_data = get_data_traversal_rules()
     traversal_contexts_collection = speckle_data.traverse(version_root_object)
 
     # Checking rules
-    is_revit_parameter = ParameterRules.speckle_type_rule("Objects.BuiltElements.Revit.Parameter")
-    has_forbidden_prefix = ParameterRules.forbidden_prefix_rule(function_inputs.forbidden_parameter_prefix)
+    is_revit_parameter = ParameterRules.speckle_type_rule(
+        "Objects.BuiltElements.Revit.Parameter"
+    )
+    has_forbidden_prefix = ParameterRules.forbidden_prefix_rule(
+        function_inputs.forbidden_parameter_prefix
+    )
 
     # Actions
     removal_action = PrefixRemovalAction(function_inputs.forbidden_parameter_prefix)
@@ -100,8 +104,9 @@ def automate_function(
     for action in [removal_action]:
         action.report(automate_context)
 
-    new_version_id = automate_context.create_new_version_in_project(version_root_object, "cleansed",
-                                                                    "Cleansed Parameters")
+    new_version_id = automate_context.create_new_version_in_project(
+        version_root_object, "cleansed", "Cleansed Parameters"
+    )
 
     if not new_version_id:
         automate_context.mark_run_failed("Failed to create a new version.")
@@ -111,32 +116,6 @@ def automate_function(
     automate_context.mark_run_success("Actions applied and reports generated.")
 
 
-def get_data_traversal() -> GraphTraversal:
-    """
-    This function is responsible for navigating through the Speckle data
-    hierarchy and providing contexts to be checked and acted upon.
-    """
-    display_value_property_aliases = {"displayValue", "@displayValue"}
-    elements_property_aliases = {"elements", "@elements"}
-
-    display_value_rule = TraversalRule(
-        [
-            lambda o: any(
-                getattr(o, alias, None) for alias in display_value_property_aliases
-            ),
-            lambda o: "Geometry" in o.speckle_type,
-        ],
-        lambda o: elements_property_aliases,
-    )
-
-    default_rule = TraversalRule(
-        [lambda _: True],
-        lambda o: o.get_member_names(),  # TODO: avoid deprecated members
-    )
-
-    return GraphTraversal([display_value_rule, default_rule])
-
-
 # make sure to call the function with the executor
 if __name__ == "__main__":
     # NOTE: always pass in the automate function by its reference, do not invoke it!