Adds dependency injection to check for root status of nodes

Author: sharpencrag

README.md

@@ -169,7 +169,10 @@ tree = ascii_tree.renderable_from_parents(
 )
 print(ascii_tree.render(tree[0]))
 ```
-Root-level nodes in your hierarchy *must* return `None` from either `parent_attr` or `parent_method`.
+By default, nodes are considered to be "root-level" if they return `None` from either `parent_attr` or `parent_method`.
+
+If your system handles root status differently, you can pass in a callback function,
+`is_root_callback`, to `renderable_from_parents` to override this behavior.  `is_root_callback` takes a single argument, the node in question, and should return `True` if the node is a root node, and `False` otherwise.
 
 Note that `renderable_from_parents` always **returns a list** -- this is to ensure that
 data with multiple root nodes are supported.
@@ -180,6 +183,7 @@ root
 └─ child
    └─ grandchild
 ```
+
 ## Rendering Styles
 
 `ascii-tree` currently provides three styles of rendering:

noxfile.py

@@ -1,7 +1,7 @@
 import nox
 
 
-@nox.session
+@nox.session(python=["3.7", "3.8", "3.9", "3.10", "3.11"])
 def tests(session):
     session.install(".")
     session.run("pytest", "tests")

src/ascii_tree/__init__.py

@@ -75,6 +75,7 @@ def renderable_from_parents(
     objs: t.Sequence[t.Any],
     parent_attr: t.Optional[str] = None,
     parent_method: t.Optional[t.Callable] = None,
+    is_root_callback: t.Optional[t.Callable[[t.Any], bool]] = None,
     display_attr: t.Optional[str] = None,
     display_method: t.Optional[t.Callable] = None,
 ) -> t.List[TextRenderNode]:
@@ -89,9 +90,9 @@ def renderable_from_parents(
     are multiple roots in your given sequence.
     """
 
-    roots = []
-
     def get_parent(obj):
+        if is_root_callback and is_root_callback(ancestor):
+            return None
         if parent_attr:
             return getattr(obj, parent_attr)
         elif parent_method:
@@ -107,7 +108,7 @@ def make_node(obj):
             children_method=lambda: list(),
         )
 
-    node_dict = dict()
+    node_dict: t.Dict[TextRenderNode, TextRenderNode] = dict()
     roots = set()
     for obj in objs:
         ancestor = obj
@@ -145,18 +146,18 @@ def renderable_dir_tree(
     """Create a TextRenderNode tree from a given file system path.
 
     Args:
-        path (str): The path to the directory to build the tree from.
-        recursive (bool, optional): Whether to build the tree recursively.
+        path: The path to the directory to build the tree from.
+        recursive: Whether to build the tree recursively.
             Defaults to True.
-        max_depth (Optional[int], optional): The maximum depth to build the tree
+        max_depth: The maximum depth to build the tree
             to. If not specified, the tree will be built to its full depth. Defaults
             to None.
-        slash_after_dir (bool, optional): Whether to add a forward slash after
+        slash_after_dir: Whether to add a forward slash after
             directories in the tree. Defaults to True.
-        ellipsis_after_max_depth (bool, optional): Whether to add an ellipsis
+        ellipsis_after_max_depth: Whether to add an ellipsis
             after the last directory in a path that reaches the maximum depth.
             Defaults to True.
-        skip_if_no_permission (bool, optional): Whether to skip adding a node to
+        skip_if_no_permission: Whether to skip adding a node to
             the tree if permission is denied to access it. Defaults to True.
 
 

tests/test_interface.py

@@ -66,6 +66,15 @@ def __init__(self, name, parent=None):
             self.name = name
             self.parent = parent
 
+    class DummyWithDifferentParentSentinel:
+        # This tests the ability to use a callback to determine if a node is root
+        def __init__(self, name, parent=None):
+            self.name = name
+            self.parent = parent or 5
+
+        def is_root(self):
+            return self.parent == 5
+
     def test_single_root(self):
         obj1 = self.DummyObject("Root")
         obj2 = self.DummyObject("Child", parent=obj1)
@@ -88,9 +97,23 @@ def test_multiple_roots(self):
             display_attr="name",
         )
         assert len(nodes) == 2
-        assert sorted([node.display for node in nodes]) == ["Root1", "Root2"]
+        nodes = sorted(nodes, key=lambda node: node.display)
+        assert [node.display for node in nodes] == ["Root1", "Root2"]
         assert nodes[1].children[0].display == "Child"
 
+    def test_root_callback(self):
+        obj = self.DummyWithDifferentParentSentinel("Root")
+        obj2 = self.DummyWithDifferentParentSentinel("Child", parent=obj)
+        nodes = renderable_from_parents(
+            [obj2],
+            parent_attr="parent",
+            display_attr="name",
+            is_root_callback=(lambda node: node.is_root()),
+        )
+        assert len(nodes) == 1
+        assert nodes[0].display == "Root"
+        assert nodes[0].children[0].display == "Child"
+
 
 if __name__ == "__main__":
     pytest.main()