Migrate more functions to plugin (#415)

* feat: nested_key_dict to tree to support child_key=None * feat: tree to nested_key_dict to support child_key=None * refactor: enable plugin - retains docstrings, help, and suggestions * docs: update CHANGELOG * refactor: migrate more methods to plugin * test: fix codecov * refactor: more clean up

Author: kayjan

bigtree/binarytree/binarytree.py

@@ -29,7 +29,7 @@ def __init__(self, root: binarynode.BinaryNode):
  # Append methods
  "from_heapq_list": construct.list_to_binarytree,
  },
- is_classmethod=True,
+ method="class",
 )
 BinaryTree.register_plugins(
  {

bigtree/dag/dag.py

@@ -2,26 +2,12 @@
 
 import copy
 import functools
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Literal, TypeVar
 
 from bigtree.dag import construct, export
 from bigtree.node import dagnode
 from bigtree.utils import iterators
 
-try:
- import pandas as pd
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pd = MagicMock()
-
-try:
- import pydot
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pydot = MagicMock()
-
 
 class DAG:
  """
@@ -41,36 +27,38 @@ def __init__(self, dag: dagnode.DAGNode):
 
  @classmethod
  def register_plugin(
- cls, name: str, func: Callable[..., Any], is_classmethod: bool
+ cls, name: str, func: Callable[..., Any], method: Literal["default", "class"]
  ) -> None:
  base_func = func.func if isinstance(func, functools.partial) else func
 
- if is_classmethod:
+ if method == "default":
+
+ def wrapper(self, *args, **kwargs): # type: ignore
+ return func(self.dag, *args, **kwargs)
+
+ else:
 
  def wrapper(cls, *args, **kwargs): # type: ignore
  construct_kwargs = {**cls.construct_kwargs, **kwargs}
  root_node = func(*args, **construct_kwargs)
  return cls(root_node)
 
- else:
-
- def wrapper(self, *args, **kwargs): # type: ignore
- return func(self.dag, *args, **kwargs)
-
  functools.update_wrapper(wrapper, base_func)
  wrapper.__name__ = name
- if is_classmethod:
+ if method == "class":
  setattr(cls, name, classmethod(wrapper)) # type: ignore
  else:
  setattr(cls, name, wrapper)
  cls._plugins[name] = func
 
  @classmethod
  def register_plugins(
- cls, mapping: dict[str, Callable[..., Any]], is_classmethod: bool = False
+ cls,
+ mapping: dict[str, Callable[..., Any]],
+ method: Literal["default", "class"] = "default",
  ) -> None:
  for name, func in mapping.items():
- cls.register_plugin(name, func, is_classmethod)
+ cls.register_plugin(name, func, method)
 
  # Magic methods
  def __getitem__(self, child_name: str) -> "DAG":
@@ -140,7 +128,7 @@ def __repr__(self) -> str:
  "from_dict": construct.dict_to_dag,
  "from_list": construct.list_to_dag,
  },
- is_classmethod=True,
+ method="class",
 )
 DAG.register_plugins(
  {

bigtree/tree/tree.py

@@ -2,57 +2,12 @@
 
 import copy
 import functools
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Literal, TypeVar
 
 from bigtree.node import basenode, binarynode, node
 from bigtree.tree import construct, export, helper, query, search
 from bigtree.utils import iterators
 
-try:
- import pandas as pd
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pd = MagicMock()
-
-try:
- import polars as pl
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pl = MagicMock()
-
-
-try:
- import pydot
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pydot = MagicMock()
-
-try:
- from PIL import Image, ImageDraw, ImageFont
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- Image = ImageDraw = ImageFont = MagicMock()
-
-try:
- import matplotlib as mpl
- from matplotlib.colors import Normalize
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- mpl = MagicMock()
- Normalize = MagicMock()
-
-try:
- import pyvis
-except ImportError: # pragma: no cover
- from unittest.mock import MagicMock
-
- pyvis = MagicMock()
-
 try:
  import matplotlib.pyplot as plt
 except ImportError: # pragma: no cover
@@ -85,36 +40,51 @@ def __init__(self, root: node.Node):
 
  @classmethod
  def register_plugin(
- cls, name: str, func: Callable[..., Any], is_classmethod: bool
+ cls,
+ name: str,
+ func: Callable[..., Any],
+ method: Literal["default", "class", "helper", "diff"],
  ) -> None:
  base_func = func.func if isinstance(func, functools.partial) else func
 
- if is_classmethod:
+ if method == "default":
+
+ def wrapper(self, *args, **kwargs): # type: ignore
+ return func(self.node, *args, **kwargs)
+
+ elif method == "class":
 
  def wrapper(cls, *args, **kwargs): # type: ignore
  construct_kwargs = {**cls.construct_kwargs, **kwargs}
  root_node = func(*args, **construct_kwargs)
  return cls(root_node)
 
- else:
+ elif method == "helper":
 
  def wrapper(self, *args, **kwargs): # type: ignore
- return func(self.node, *args, **kwargs)
+ return type(self)(func(self.node, *args, **kwargs))
+
+ else:
+
+ def wrapper(self, other_tree: T, *args, **kwargs): # type: ignore
+ return func(self.node, other_tree.node, *args, **kwargs)
 
  functools.update_wrapper(wrapper, base_func)
  wrapper.__name__ = name
- if is_classmethod:
+ if method == "class":
  setattr(cls, name, classmethod(wrapper)) # type: ignore
  else:
  setattr(cls, name, wrapper)
  cls._plugins[name] = func
 
  @classmethod
  def register_plugins(
- cls, mapping: dict[str, Callable[..., Any]], is_classmethod: bool = False
+ cls,
+ mapping: dict[str, Callable[..., Any]],
+ method: Literal["default", "class", "helper", "diff"] = "default",
  ) -> None:
  for name, func in mapping.items():
- cls.register_plugin(name, func, is_classmethod)
+ cls.register_plugin(name, func, method)
 
  def show(self, **kwargs: Any) -> None:
  self.node.show(**kwargs)
@@ -143,38 +113,6 @@ def depth(self) -> int:
  """
  return self.node.max_depth
 
- # Helper methods
- def clone(self, node_type: type[BaseNodeT]) -> "Tree":
- """See `clone_tree` for full details.
-
- Accepts the same arguments as `clone_tree`.
- """
- return type(self)(helper.clone_tree(self.node, node_type)) # type: ignore
-
- def prune(self, *args: Any, **kwargs: Any) -> "Tree":
- """See `prune_tree` for full details.
-
- Accepts the same arguments as `prune_tree`.
- """
- return type(self)(helper.prune_tree(self.node, *args, **kwargs))
-
- def diff_dataframe(self, other_tree: T, *args: Any, **kwargs: Any) -> pd.DataFrame:
- """See `get_tree_diff_dataframe` for full details.
-
- Accepts the same arguments as `get_tree_diff_dataframe`.
- """
- return helper.get_tree_diff_dataframe(
- self.node, other_tree.node, *args, **kwargs
- )
-
- def diff(self, other_tree: T, *args: Any, **kwargs: Any) -> node.Node:
- """See `get_tree_diff` for full details.
-
- Accepts the same arguments as `get_tree_diff`.
- """
- return helper.get_tree_diff(self.node, other_tree.node, *args, **kwargs)
-
- # Plot methods
  def plot(self, *args: Any, **kwargs: Any) -> plt.Figure:
  """Plot tree in line form. Accepts args and kwargs for matplotlib.pyplot.plot() function.
 
@@ -259,7 +197,7 @@ def __repr__(self) -> str:
  "from_str": construct.str_to_tree,
  "from_newick": construct.newick_to_tree,
  },
- is_classmethod=True,
+ method="class",
 )
 
 Tree.register_plugins(
@@ -309,3 +247,19 @@ def __repr__(self) -> str:
  "zigzaggroup_iter": iterators.zigzaggroup_iter,
  }
 )
+Tree.register_plugins(
+ {
+ # Helper methods
+ "clone": helper.clone_tree,
+ "prune": helper.prune_tree,
+ },
+ method="helper",
+)
+Tree.register_plugins(
+ {
+ # Helper methods
+ "diff_dataframe": helper.get_tree_diff_dataframe,
+ "diff": helper.get_tree_diff,
+ },
+ method="diff",
+)