Merge pull request #291 from kayjan/feat/plot-tree

Add plotting for tree

CHANGELOG.md

@@ -6,6 +6,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.21.0] - TBD
+### Added:
+- Tree Plot: Plot tree using matplotlib library, added matplotlib as optional dependency.
+- BaseNode: Add plot method.
+### Changed:
+- Misc: Optional dependencies imported as MagicMock
+
 ## [0.20.1] - 2024-08-24
 ### Changed:
 - Misc: Documentation update contributing instructions.
@@ -638,7 +645,8 @@ ignore null attribute columns.
 - Utility Iterator: Tree traversal methods.
 - Workflow To Do App: Tree use case with to-do list implementation.
 
-[Unreleased]: https://github.com/kayjan/bigtree/compare/0.20.1...HEAD
+[Unreleased]: https://github.com/kayjan/bigtree/compare/0.21.0...HEAD
+[0.21.0]: https://github.com/kayjan/bigtree/compare/0.20.1...0.21.0
 [0.20.1]: https://github.com/kayjan/bigtree/compare/0.20.0...0.20.1
 [0.20.0]: https://github.com/kayjan/bigtree/compare/0.19.4...0.20.0
 [0.19.4]: https://github.com/kayjan/bigtree/compare/0.19.3...0.19.4

README.md

@@ -67,6 +67,7 @@ For **Tree** implementation, there are 9 main components.
    4. Get difference between two trees
 7. [**📊 Plotting Tree**](https://bigtree.readthedocs.io/en/stable/bigtree/utils/plot/)
    1. Enhanced Reingold Tilford Algorithm to retrieve (x, y) coordinates for a tree structure
+   2. Plot tree using matplotlib (optional dependency)
 8. [**🔨 Exporting Tree**](https://bigtree.readthedocs.io/en/stable/bigtree/tree/export/)
    1. Print to console, in vertical or horizontal orientation
    2. Export to *Newick string notation*, *dictionary*, *nested dictionary*, *pandas DataFrame*, or *polars DataFrame*

bigtree/__init__.py

@@ -93,6 +93,6 @@
     zigzag_iter,
     zigzaggroup_iter,
 )
-from bigtree.utils.plot import reingold_tilford
+from bigtree.utils.plot import plot_tree, reingold_tilford
 from bigtree.workflows.app_calendar import Calendar
 from bigtree.workflows.app_todo import AppToDo

bigtree/dag/construct.py

@@ -16,7 +16,9 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pd = MagicMock()
 
 __all__ = ["list_to_dag", "dict_to_dag", "dataframe_to_dag"]
 

bigtree/dag/export.py

@@ -13,12 +13,16 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pd = MagicMock()
 
 try:
     import pydot
 except ImportError:  # pragma: no cover
-    pydot = None
+    from unittest.mock import MagicMock
+
+    pydot = MagicMock()
 
 __all__ = ["dag_to_list", "dag_to_dict", "dag_to_dataframe", "dag_to_dot"]
 

bigtree/node/basenode.py

@@ -8,6 +8,11 @@
 from bigtree.utils.exceptions import CorruptedTreeError, LoopError, TreeError
 from bigtree.utils.iterators import preorder_iter
 
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # pragma: no cover
+    plt = None
+
 
 class BaseNode:
     """
@@ -115,6 +120,7 @@ class BaseNode:
     6. ``extend(nodes: List[Self])``: Add multiple children to node
     7. ``copy()``: Deep copy self
     8. ``sort()``: Sort child nodes
+    9. ``plot()``: Plot tree in line form
 
     ----
 
@@ -727,6 +733,7 @@ def copy(self: T) -> T:
 
     def sort(self: T, **kwargs: Any) -> None:
         """Sort children, possible keyword arguments include ``key=lambda node: node.name``, ``reverse=True``
+        Accepts kwargs for sort() function.
 
         Examples:
             >>> from bigtree import Node, print_tree
@@ -747,6 +754,24 @@ def sort(self: T, **kwargs: Any) -> None:
         children.sort(**kwargs)
         self.__children = children
 
+    def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
+        """Plot tree in line form.
+        Accepts args and kwargs for matplotlib.pyplot.plot() function.
+
+        Examples:
+            >>> import matplotlib.pyplot as plt
+            >>> from bigtree import list_to_tree
+            >>> path_list = ["a/b/d", "a/b/e/g", "a/b/e/h", "a/c/f"]
+            >>> root = list_to_tree(path_list)
+            >>> root.plot("-ok")
+            <Figure size 1280x960 with 1 Axes>
+        """
+        from bigtree.utils.plot import plot_tree, reingold_tilford
+
+        if self.get_attr("x") is None or self.get_attr("y") is None:
+            reingold_tilford(self)
+        return plot_tree(self, *args, **kwargs)
+
     def __copy__(self: T) -> T:
         """Shallow copy self
 

bigtree/tree/construct.py

@@ -24,12 +24,16 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pd = MagicMock()
 
 try:
     import polars as pl
 except ImportError:  # pragma: no cover
-    pl = None
+    from unittest.mock import MagicMock
+
+    pl = MagicMock()
 
 __all__ = [
     "add_path_to_tree",

bigtree/tree/export.py

@@ -28,22 +28,30 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pd = MagicMock()
 
 try:
     import polars as pl
 except ImportError:  # pragma: no cover
-    pl = None
+    from unittest.mock import MagicMock
+
+    pl = MagicMock()
 
 try:
     import pydot
 except ImportError:  # pragma: no cover
-    pydot = None
+    from unittest.mock import MagicMock
+
+    pydot = MagicMock()
 
 try:
     from PIL import Image, ImageDraw, ImageFont
 except ImportError:  # pragma: no cover
-    Image = ImageDraw = ImageFont = None
+    from unittest.mock import MagicMock
+
+    Image = ImageDraw = ImageFont = MagicMock()
 
 
 __all__ = [
@@ -73,9 +81,10 @@ def print_tree(
     attr_omit_null: bool = False,
     attr_bracket: List[str] = ["[", "]"],
     style: Union[str, Iterable[str], BasePrintStyle] = "const",
-    **print_kwargs: Any,
+    **kwargs: Any,
 ) -> None:
     """Print tree to console, starting from `tree`.
+    Accepts kwargs for print() function.
 
     - Able to select which node to print from, resulting in a subtree, using `node_name_or_path`
     - Able to customize for maximum depth to print, using `max_depth`
@@ -91,8 +100,6 @@ def print_tree(
     - (BasePrintStyle): `ANSIPrintStyle`, `ASCIIPrintStyle`, `ConstPrintStyle`, `ConstBoldPrintStyle`, `RoundedPrintStyle`,
     `DoublePrintStyle` style or inherit from `BasePrintStyle`
 
-    Remaining kwargs are passed without modification to python's `print` function.
-
     Examples:
         **Printing tree**
 
@@ -249,7 +256,7 @@ def print_tree(
             if attr_str:
                 attr_str = f" {attr_bracket_open}{attr_str}{attr_bracket_close}"
         node_str = f"{_node.node_name}{attr_str}"
-        print(f"{pre_str}{fill_str}{node_str}", **print_kwargs)
+        print(f"{pre_str}{fill_str}{node_str}", **kwargs)
 
 
 def yield_tree(
@@ -433,9 +440,10 @@ def hprint_tree(
     max_depth: int = 0,
     intermediate_node_name: bool = True,
     style: Union[str, Iterable[str], BaseHPrintStyle] = "const",
-    **print_kwargs: Any,
+    **kwargs: Any,
 ) -> None:
     """Print tree in horizontal orientation to console, starting from `tree`.
+    Accepts kwargs for print() function.
 
     - Able to select which node to print from, resulting in a subtree, using `node_name_or_path`
     - Able to customize for maximum depth to print, using `max_depth`
@@ -448,8 +456,6 @@ def hprint_tree(
     - (BaseHPrintStyle): `ANSIHPrintStyle`, `ASCIIHPrintStyle`, `ConstHPrintStyle`, `ConstBoldHPrintStyle`,
     `RoundedHPrintStyle`, `DoubleHPrintStyle` style or inherit from BaseHPrintStyle
 
-    Remaining kwargs are passed without modification to python's `print` function.
-
     Examples:
         **Printing tree**
 
@@ -549,7 +555,7 @@ def hprint_tree(
         max_depth=max_depth,
         style=style,
     )
-    print("\n".join(result), **print_kwargs)
+    print("\n".join(result), **kwargs)
 
 
 def hyield_tree(

bigtree/utils/assertions.py

@@ -5,12 +5,16 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pd = MagicMock()
 
 try:
     import polars as pl
 except ImportError:  # pragma: no cover
-    pl = None
+    from unittest.mock import MagicMock
+
+    pl = MagicMock()
 
 
 if TYPE_CHECKING:

bigtree/utils/exceptions.py

@@ -114,6 +114,28 @@ def wrapper(*args: Any, **kwargs: Any) -> T:
     return wrapper
 
 
+def optional_dependencies_matplotlib(
+    func: Callable[..., T]
+) -> Callable[..., T]:  # pragma: no cover
+    """
+    This is a decorator which can be used to import optional matplotlib dependency.
+    It will raise a ImportError if the module is not found.
+    """
+
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> T:
+        try:
+            import matplotlib.pyplot as plt  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "matplotlib not available. Please perform a\n\n"
+                "pip install 'bigtree[matplotlib]'\n\nto install required dependencies"
+            ) from None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 def optional_dependencies_image(
     package_name: str = "",
 ) -> Callable[[Callable[..., T]], Callable[..., T]]:

bigtree/utils/plot.py

@@ -1,9 +1,20 @@
-from typing import Optional, TypeVar
+from typing import Any, Optional, TypeVar
 
 from bigtree.node.basenode import BaseNode
+from bigtree.utils.exceptions import optional_dependencies_matplotlib
+from bigtree.utils.iterators import preorder_iter
+
+try:
+    import matplotlib.pyplot as plt
+except ImportError:  # pragma: no cover
+    from unittest.mock import MagicMock
+
+    plt = MagicMock()
+
 
 __all__ = [
     "reingold_tilford",
+    "plot_tree",
 ]
 
 T = TypeVar("T", bound=BaseNode)
@@ -73,6 +84,49 @@ def reingold_tilford(
     _third_pass(tree_node, x_adjustment)
 
 
+@optional_dependencies_matplotlib
+def plot_tree(
+    tree_node: T, *args: Any, ax: Optional[plt.Axes] = None, **kwargs: Any
+) -> plt.Figure:
+    """Plot tree in line form. Tree should have `x` and `y` attribute from Reingold Tilford.
+    Accepts existing matplotlib Axes. Accepts args and kwargs for matplotlib.pyplot.plot() function.
+
+    Examples:
+        >>> import matplotlib.pyplot as plt
+        >>> from bigtree import list_to_tree, plot_tree, reingold_tilford
+        >>> path_list = ["a/b/d", "a/b/e/g", "a/b/e/h", "a/c/f"]
+        >>> root = list_to_tree(path_list)
+        >>> reingold_tilford(root)
+        >>> plot_tree(root, "-ok")
+        <Figure size 1280x960 with 1 Axes>
+
+    Args:
+        tree_node (BaseNode): tree to plot
+        ax (plt.Axes): axes to add Figure to
+    """
+    if ax:
+        fig = ax.get_figure()
+    else:
+        fig = plt.figure()
+        ax = fig.add_subplot(111)
+
+    for node in preorder_iter(tree_node):
+        if not node.is_root:
+            try:
+                ax.plot(
+                    [node.x, node.parent.x],  # type: ignore
+                    [node.y, node.parent.y],  # type: ignore
+                    *args,
+                    **kwargs,
+                )
+            except AttributeError:
+                raise RuntimeError(
+                    "No x or y coordinates detected. "
+                    "Please run reingold_tilford algorithm to retrieve coordinates."
+                )
+    return fig
+
+
 def _first_pass(
     tree_node: T, sibling_separation: float, subtree_separation: float
 ) -> None:

bigtree/workflows/app_calendar.py

@@ -11,7 +11,9 @@
 try:
     import pandas as pd
 except ImportError:  # pragma: no cover
-    pd = None
+    from unittest.mock import MagicMock
+
+    pydot = MagicMock()
 
 
 class Calendar:

docs/home/tree.md

@@ -50,6 +50,7 @@ For **Tree** implementation, there are 9 main components.
 
 ## [**📊 Plotting Tree**](../bigtree/utils/plot.md)
 - Enhanced Reingold Tilford Algorithm to retrieve (x, y) coordinates for a tree structure
+- Plot tree using matplotlib (optional dependency)
 
 ## [**🔨 Exporting Tree**](../bigtree/tree/export.md)
 - Print to console, in vertical or horizontal orientation

pyproject.toml

@@ -37,6 +37,7 @@ Source = "https://github.com/kayjan/bigtree"
 
 [project.optional-dependencies]
 all = [
+  "matplotlib",
   "pandas",
   "polars",
   "pydot",
@@ -46,6 +47,7 @@ image = [
   "pydot",
   "Pillow",
 ]
+matplotlib = ["matplotlib"]
 pandas = ["pandas"]
 polars = ["polars"]
 
@@ -56,6 +58,7 @@ path = "bigtree/__init__.py"
 dependencies = [
   "black",
   "coverage",
+  "matplotlib",
   "mypy",
   "pandas",
   "Pillow",