Merge pull request #252 from kayjan/add-dataclass

Add dataclass

Author: kayjan

CHANGELOG.md

@@ -5,6 +5,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Changed:
+- Tree Exporter: Print functions to accept custom style that is implemented as dataclass, this is a more
+object-oriented way of parsing arguments.
+This affects functions `print_tree`, `yield_tree`, `hprint_tree`, and `hyield_tree`.
+The argument `custom_style` is deprecated, and argument `style` is used instead.
+**This might not be backwards-compatible!**
+- Misc: Updated docstrings to be more comprehensive for tree constructor and tree exporter.
 
 ## [0.18.3] - 2024-06-05
 ### Changed:

bigtree/__init__.py

@@ -66,6 +66,22 @@
     find_relative_paths,
     findall,
 )
+from bigtree.utils.constants import (
+    ANSIHPrintStyle,
+    ANSIPrintStyle,
+    ASCIIHPrintStyle,
+    ASCIIPrintStyle,
+    BaseHPrintStyle,
+    BasePrintStyle,
+    ConstBoldHPrintStyle,
+    ConstBoldPrintStyle,
+    ConstHPrintStyle,
+    ConstPrintStyle,
+    DoubleHPrintStyle,
+    DoublePrintStyle,
+    RoundedHPrintStyle,
+    RoundedPrintStyle,
+)
 from bigtree.utils.groot import speak_like_groot, whoami
 from bigtree.utils.iterators import (
     dag_iterator,

bigtree/tree/construct.py

@@ -1411,7 +1411,7 @@ def newick_to_tree(
     In the Newick Notation (or New Hampshire Notation)
 
     - Tree is represented in round brackets i.e., `(child1,child2,child3)parent`.
-    - If there are nested tree, they will be in nested round brackets i.e., `((grandchild1)child1,(grandchild2,grandchild3)child2)parent`.
+    - If there are nested trees, they will be in nested round brackets i.e., `((grandchild1)child1,(grandchild2,grandchild3)child2)parent`.
     - If there is length attribute, they will be beside the name i.e., `(child1:0.5,child2:0.1)parent`.
     - If there are other attributes, attributes are represented in square brackets i.e., `(child1:0.5[S:human],child2:0.1[S:human])parent[S:parent]`.
 

bigtree/tree/export.py

@@ -11,7 +11,13 @@
     assert_tree_type,
     isnull,
 )
-from bigtree.utils.constants import ExportConstants, MermaidConstants, NewickCharacter
+from bigtree.utils.constants import (
+    BaseHPrintStyle,
+    BasePrintStyle,
+    ExportConstants,
+    MermaidConstants,
+    NewickCharacter,
+)
 from bigtree.utils.exceptions import (
     optional_dependencies_image,
     optional_dependencies_pandas,
@@ -66,8 +72,7 @@ def print_tree(
     attr_list: Iterable[str] = [],
     attr_omit_null: bool = False,
     attr_bracket: List[str] = ["[", "]"],
-    style: str = "const",
-    custom_style: Iterable[str] = [],
+    style: Union[str, Iterable[str], BasePrintStyle] = "const",
 ) -> None:
     """Print tree to console, starting from `tree`.
 
@@ -76,10 +81,14 @@ def print_tree(
     - Able to choose which attributes to show or show all attributes, using `attr_name_filter` and `all_attrs`
     - Able to omit showing of attributes if it is null, using `attr_omit_null`
     - Able to customize open and close brackets if attributes are shown, using `attr_bracket`
-    - Able to customize style, to choose from `ansi`, `ascii`, `const`, `const_bold`, `rounded`, `double`, and `custom` style
-        - Default style is `const` style
-        - If style is set to custom, user can choose their own style for stem, branch and final stem icons
-        - Stem, branch, and final stem symbol should have the same number of characters
+    - Able to customize style, to choose from str, List[str], or inherit from BasePrintStyle, using `style`
+
+    For style,
+
+    - (str): `ansi`, `ascii`, `const` (default), `const_bold`, `rounded`, `double`  style
+    - (List[str]): Choose own style for stem, branch, and final stem icons, they must have the same number of characters
+    - (BasePrintStyle): `ANSIPrintStyle`, `ASCIIPrintStyle`, `ConstPrintStyle`, `ConstBoldPrintStyle`, `RoundedPrintStyle`,
+    `DoublePrintStyle` style or inherit from `BasePrintStyle`
 
     Examples:
         **Printing tree**
@@ -169,6 +178,16 @@ def print_tree(
         ║   ╚══ e
         ╚══ c
 
+        **Custom Styles**
+
+        >>> from bigtree import ANSIPrintStyle
+        >>> print_tree(root, style=ANSIPrintStyle)
+        a
+        |-- b
+        |   |-- d
+        |   `-- e
+        `-- c
+
     Args:
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing
@@ -177,15 +196,13 @@ def print_tree(
         attr_list (Iterable[str]): list of node attributes to print, optional
         attr_omit_null (bool): indicator whether to omit showing of null attributes, defaults to False
         attr_bracket (List[str]): open and close bracket for `all_attrs` or `attr_list`
-        style (str): style of print, defaults to const style
-        custom_style (Iterable[str]): style of stem, branch and final stem, used when `style` is set to 'custom'
+        style (Union[str, Iterable[str], BasePrintStyle]): style of print, defaults to const
     """
     for pre_str, fill_str, _node in yield_tree(
         tree=tree,
         node_name_or_path=node_name_or_path,
         max_depth=max_depth,
         style=style,
-        custom_style=custom_style,
     ):
         # Get node_str (node name and attributes)
         attr_str = ""
@@ -222,17 +239,20 @@ def yield_tree(
     tree: T,
     node_name_or_path: str = "",
     max_depth: int = 0,
-    style: str = "const",
-    custom_style: Iterable[str] = [],
+    style: Union[str, Iterable[str], BasePrintStyle] = "const",
 ) -> Iterable[Tuple[str, str, T]]:
     """Generator method for customizing printing of tree, starting from `tree`.
 
     - 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`
-    - Able to customize style, to choose from `ansi`, `ascii`, `const`, `const_bold`, `rounded`, `double`, and `custom` style
-        - Default style is `const` style
-        - If style is set to custom, user can choose their own style for stem, branch and final stem icons
-        - Stem, branch, and final stem symbol should have the same number of characters
+    - Able to customize style, to choose from str, List[str], or inherit from BasePrintStyle, using `style`
+
+    For style,
+
+    - (str): `ansi`, `ascii`, `const` (default), `const_bold`, `rounded`, `double`  style
+    - (List[str]): Choose own style for stem, branch, and final stem icons, they must have the same number of characters
+    - (BasePrintStyle): `ANSIPrintStyle`, `ASCIIPrintStyle`, `ConstPrintStyle`, `ConstBoldPrintStyle`, `RoundedPrintStyle`,
+    `DoublePrintStyle` style or inherit from `BasePrintStyle`
 
     Examples:
         **Yield tree**
@@ -315,6 +335,17 @@ def yield_tree(
         ║   ╚══ e
         ╚══ c
 
+        **Custom Styles**
+
+        >>> from bigtree import ANSIPrintStyle
+        >>> for branch, stem, node in yield_tree(root, style=ANSIPrintStyle):
+        ...     print(f"{branch}{stem}{node.node_name}")
+        a
+        |-- b
+        |   |-- d
+        |   `-- e
+        `-- c
+
         **Printing Attributes**
 
         >>> for branch, stem, node in yield_tree(root, style="const"):
@@ -329,30 +360,26 @@ def yield_tree(
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing, optional
         max_depth (int): maximum depth of tree to print, based on `depth` attribute, optional
-        style (str): style of print, defaults to const
-        custom_style (Iterable[str]): style of stem, branch and final stem, used when `style` is set to 'custom'
+        style (Union[str, Iterable[str], BasePrintStyle]): style of print, defaults to const
     """
     from bigtree.tree.helper import get_subtree
 
-    available_styles = ExportConstants.PRINT_STYLES
-    assert_style_in_dict(style, available_styles)
-
     tree = get_subtree(tree, node_name_or_path, max_depth)
 
     # Set style
-    if style == "custom":
-        if len(list(custom_style)) != 3:
-            raise ValueError(
-                "Custom style selected, please specify the style of stem, branch, and final stem in `custom_style`"
-            )
-        style_stem, style_branch, style_stem_final = custom_style
-    else:
+    if isinstance(style, str):
+        available_styles = ExportConstants.PRINT_STYLES
+        assert_style_in_dict(style, available_styles)
         style_stem, style_branch, style_stem_final = available_styles[style]
-
-    if not len(style_stem) == len(style_branch) == len(style_stem_final):
+    elif isinstance(style, list) and len(list(style)) != 3:
         raise ValueError(
-            "`style_stem`, `style_branch`, and `style_stem_final` are of different length"
+            "Please specify the style of stem, branch, and final stem in `style`"
         )
+    else:
+        style_stem, style_branch, style_stem_final = style  # type: ignore[misc]
+
+    if not len(style_stem) == len(style_branch) == len(style_stem_final):
+        raise ValueError("`stem`, `branch`, and `stem_final` are of different length")
 
     gap_str = " " * len(style_stem)
     unclosed_depth = set()
@@ -388,17 +415,20 @@ def hprint_tree(
     node_name_or_path: str = "",
     max_depth: int = 0,
     intermediate_node_name: bool = True,
-    style: str = "const",
-    custom_style: Iterable[str] = [],
+    style: Union[str, Iterable[str], BaseHPrintStyle] = "const",
 ) -> None:
     """Print tree in horizontal orientation to console, starting from `tree`.
 
     - 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`
-    - Able to customize style, to choose from `ansi`, `ascii`, `const`, `const_bold`, `rounded`, `double`, and `custom` style
-        - Default style is `const` style
-        - If style is set to custom, user can choose their own style icons
-        - Style icons should have the same number of characters
+    - Able to customize style, to choose from str, List[str], or inherit from BaseHPrintStyle, using `style`
+
+    For style,
+
+    - (str): `ansi`, `ascii`, `const` (default), `const_bold`, `rounded`, `double`  style
+    - (List[str]): Choose own style icons, they must have the same number of characters
+    - (BaseHPrintStyle): `ANSIHPrintStyle`, `ASCIIHPrintStyle`, `ConstHPrintStyle`, `ConstBoldHPrintStyle`,
+    `RoundedHPrintStyle`, `DoubleHPrintStyle` style or inherit from BaseHPrintStyle
 
     Examples:
         **Printing tree**
@@ -465,21 +495,28 @@ def hprint_tree(
         ═ a ═╣     ╚═ e
              ╚═ c
 
+        **Custom Styles**
+
+        >>> from bigtree import ANSIHPrintStyle
+        >>> hprint_tree(root, style=ANSIHPrintStyle)
+                   /- d
+             /- b -+
+        - a -+     \\- e
+             \\- c
+
     Args:
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing
         max_depth (int): maximum depth of tree to print, based on `depth` attribute, optional
         intermediate_node_name (bool): indicator if intermediate nodes have node names, defaults to True
-        style (str): style of print, defaults to const style
-        custom_style (Iterable[str]): style of icons, used when `style` is set to 'custom'
+        style (Union[str, Iterable[str], BaseHPrintStyle]): style of print, defaults to const
     """
     result = hyield_tree(
         tree,
         node_name_or_path=node_name_or_path,
         intermediate_node_name=intermediate_node_name,
         max_depth=max_depth,
         style=style,
-        custom_style=custom_style,
     )
     print("\n".join(result))
 
@@ -489,17 +526,20 @@ def hyield_tree(
     node_name_or_path: str = "",
     max_depth: int = 0,
     intermediate_node_name: bool = True,
-    style: str = "const",
-    custom_style: Iterable[str] = [],
+    style: Union[str, Iterable[str], BaseHPrintStyle] = "const",
 ) -> List[str]:
     """Yield tree in horizontal orientation to console, starting from `tree`.
 
     - 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`
-    - Able to customize style, to choose from `ansi`, `ascii`, `const`, `const_bold`, `rounded`, `double`, and `custom` style
-        - Default style is `const` style
-        - If style is set to custom, user can choose their own style icons
-        - Style icons should have the same number of characters
+    - Able to customize style, to choose from str, List[str], or inherit from BaseHPrintStyle, using `style`
+
+    For style,
+
+    - (str): `ansi`, `ascii`, `const` (default), `const_bold`, `rounded`, `double`  style
+    - (List[str]): Choose own style icons, they must have the same number of characters
+    - (BaseHPrintStyle): `ANSIHPrintStyle`, `ASCIIHPrintStyle`, `ConstHPrintStyle`, `ConstBoldHPrintStyle`,
+    `RoundedHPrintStyle`, `DoubleHPrintStyle` style or inherit from BaseHPrintStyle
 
     Examples:
         **Printing tree**
@@ -567,13 +607,21 @@ def hyield_tree(
         ═ a ═╣     ╚═ e
              ╚═ c
 
+        **Custom Styles**
+
+        >>> from bigtree import ANSIHPrintStyle
+        >>> hprint_tree(root, style=ANSIHPrintStyle)
+                   /- d
+             /- b -+
+        - a -+     \\- e
+             \\- c
+
     Args:
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing
         max_depth (int): maximum depth of tree to print, based on `depth` attribute, optional
         intermediate_node_name (bool): indicator if intermediate nodes have node names, defaults to True
-        style (str): style of print, defaults to const style
-        custom_style (Iterable[str]): style of icons, used when `style` is set to 'custom'
+        style (Union[str, Iterable[str], BaseHPrintStyle]): style of print, defaults to const
 
     Returns:
         (List[str])
@@ -582,17 +630,12 @@ def hyield_tree(
 
     from bigtree.tree.helper import get_subtree
 
-    available_styles = ExportConstants.HPRINT_STYLES
-    assert_style_in_dict(style, available_styles)
-
     tree = get_subtree(tree, node_name_or_path, max_depth)
 
     # Set style
-    if style == "custom":
-        if len(list(custom_style)) != 7:
-            raise ValueError(
-                "Custom style selected, please specify the style of 7 icons in `custom_style`"
-            )
+    if isinstance(style, str):
+        available_styles = ExportConstants.HPRINT_STYLES
+        assert_style_in_dict(style, available_styles)
         (
             style_first_child,
             style_subsequent_child,
@@ -601,7 +644,9 @@ def hyield_tree(
             style_last_child,
             style_stem,
             style_branch,
-        ) = custom_style
+        ) = available_styles[style]
+    elif isinstance(style, list) and len(list(style)) != 7:
+        raise ValueError("Please specify the style of 7 icons in `style`")
     else:
         (
             style_first_child,
@@ -611,7 +656,8 @@ def hyield_tree(
             style_last_child,
             style_stem,
             style_branch,
-        ) = available_styles[style]
+        ) = style  # type: ignore[misc]
+
     if (
         not len(style_first_child)
         == len(style_subsequent_child)
@@ -622,7 +668,7 @@ def hyield_tree(
         == len(style_branch)
         == 1
     ):
-        raise ValueError("For custom style, all style icons must have length 1")
+        raise ValueError("All style icons must have length 1")
 
     # Calculate padding
     space = " "
@@ -1673,15 +1719,15 @@ def tree_to_newick(
 
     In the Newick Notation (or New Hampshire Notation),
       - Tree is represented in round brackets i.e., `(child1,child2,child3)parent`.
-      - If there are nested tree, they will be in nested round brackets i.e., `((grandchild1)child1,(grandchild2,grandchild3)child2)parent`.
+      - If there are nested trees, they will be in nested round brackets i.e., `((grandchild1)child1,(grandchild2,grandchild3)child2)parent`.
       - If there is length attribute, they will be beside the name i.e., `(child1:0.5,child2:0.1)parent`.
       - If there are other attributes, attributes are represented in square brackets i.e., `(child1:0.5[S:human],child2:0.1[S:human])parent[S:parent]`.
 
     Customizations include:
       - Omitting names of root and intermediate nodes, default all node names are shown.
-      - Changing length separator to other symbol, default is `:`.
+      - Changing length separator to another symbol, default is `:`.
       - Adding an attribute prefix, default is `&&NHX:`.
-      - Changing the attribute separator to other symbol, default is `:`.
+      - Changing the attribute separator to another symbol, default is `:`.
 
     Examples:
         >>> from bigtree import Node, tree_to_newick