diff --git a/src/specklepy/objects/graph_traversal/default_traversal.py b/src/specklepy/objects/graph_traversal/default_traversal.py new file mode 100644 index 00000000..8c02dfb0 --- /dev/null +++ b/src/specklepy/objects/graph_traversal/default_traversal.py @@ -0,0 +1,28 @@ +from specklepy.objects.base import Base +from specklepy.objects.graph_traversal.traversal import GraphTraversal, TraversalRule + +DISPLAY_VALUE_PROPERTY_ALIASES = {"displayValue", "@displayValue"} +ELEMENTS_PROPERTY_ALIASES = {"elements", "@elements"} + + +def has_display_value(x: Base): + return any(hasattr(x, alias) for alias in DISPLAY_VALUE_PROPERTY_ALIASES) + + +def create_default_traversal_function() -> GraphTraversal: + """ + Traversal func for traversing the root object of a Speckle Model + """ + + convertible_rule = TraversalRule( + [lambda b: b.speckle_type != "Base", has_display_value], + lambda _: ELEMENTS_PROPERTY_ALIASES, + ) + + default_rule = TraversalRule( + [lambda _: True], + lambda o: o.get_member_names(), # NOTE: Unlike the C# implementation, this does not ignore Obsolete members + False, + ) + + return GraphTraversal([convertible_rule, default_rule]) diff --git a/src/specklepy/objects/graph_traversal/traversal.py b/src/specklepy/objects/graph_traversal/traversal.py index b32d0536..6ab16a7a 100644 --- a/src/specklepy/objects/graph_traversal/traversal.py +++ b/src/specklepy/objects/graph_traversal/traversal.py @@ -7,6 +7,11 @@ class ITraversalRule(Protocol): + + @property + def should_return(self) -> bool: + pass + def get_members_to_traverse(self, o: Base) -> Set[str]: """Get the members to traverse.""" pass @@ -50,10 +55,13 @@ def traverse(self, root: Base) -> Iterator[TraversalContext]: while len(stack) > 0: head = stack.pop() - yield head current = head.current active_rule = self._get_active_rule_or_default_rule(current) + + if active_rule.should_return: + yield head + members_to_traverse = active_rule.get_members_to_traverse(current) for child_prop in members_to_traverse: try: @@ -114,6 +122,11 @@ def _get_active_rule(self, o: Base) -> Optional[ITraversalRule]: class TraversalRule: _conditions: Collection[Callable[[Base], bool]] _members_to_traverse: Callable[[Base], Iterable[str]] + _should_return_to_output: bool = True + + @property + def should_return(self) -> bool: + return self._should_return_to_output def get_members_to_traverse(self, o: Base) -> Set[str]: return set(self._members_to_traverse(o))