swcgeom 0.18.3__py3-none-any.whl → 0.19.0__py3-none-any.whl

swcgeom/core/population.py

@@ -20,7 +20,7 @@ import warnings
 from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import ProcessPoolExecutor
 from functools import reduce
-from typing import Any, Optional, Protocol, TypeVar, cast, overload
+from typing import Any, Protocol, TypeVar, cast, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -39,10 +39,8 @@ T = TypeVar("T")
 class Trees(Protocol):
     """Trees protocol support index and len."""
 
-    # fmt: off
     def __getitem__(self, key: int, /) -> Tree: ...
     def __len__(self) -> int: ...
-    # fmt: on
 
 
 class LazyLoadingTrees:
@@ -54,11 +52,9 @@ class LazyLoadingTrees:
 
     def __init__(self, swcs: Iterable[str], **kwargs) -> None:
         """
-        Paramters
-        ---------
-        swcs : List of str
-        kwargs : dict[str, Any]
-            Forwarding to `Tree.from_swc`
+        Args:
+            swcs: List of str
+            kwargs: Forwarding to `Tree.from_swc`
         """
 
         super().__init__()
@@ -130,43 +126,40 @@ class Population:
 
     trees: Trees
 
-    # fmt: off
     @overload
-    def __init__(self, swcs: Iterable[str], lazy_loading: bool = ..., root: str = ..., **kwargs) -> None: ...
+    def __init__(
+        self, swcs: Iterable[str], lazy_loading: bool = ..., root: str = ..., **kwargs
+    ) -> None: ...
     @overload
     def __init__(self, trees: Trees, /, *, root: str = "") -> None: ...
-    # fmt: on
-
     def __init__(self, swcs, lazy_loading=True, root="", **kwargs) -> None:
         super().__init__()
         if len(swcs) > 0 and isinstance(swcs[0], str):
             warnings.warn(
                 "`Population(swcs)` has been replaced by "
-                "`Population(LazyLoadingTrees(swcs))` since v0.8.0 "
-                "thus we can create a population from a group of trees, "
-                " and this will be removed in next version",
+                "`Population(LazyLoadingTrees(swcs))` since v0.8.0 thus we can create "
+                "a population from a group of trees, and this will be removed in next "
+                "version",
                 DeprecationWarning,
             )
 
-            trees = LazyLoadingTrees(swcs, **kwargs)  # type: ignore
+            trees = LazyLoadingTrees(swcs, **kwargs)
             if not lazy_loading:
                 for i in range(len(swcs)):
                     trees.load(i)
         else:
             trees = swcs
 
-        self.trees = trees  # type: ignore
+        self.trees = trees
         self.root = root
 
         if len(swcs) == 0:
             warnings.warn(f"no trees in population from '{root}'")
 
-    # fmt:off
     @overload
     def __getitem__(self, key: slice) -> Trees: ...
     @overload
     def __getitem__(self, key: int) -> Tree: ...
-    # fmt:on
     def __getitem__(self, key: int | slice):
         if isinstance(key, slice):
             trees = NestTrees(self.trees, range(*key.indices(len(self))))
@@ -190,15 +183,14 @@ class Population:
         self,
         fn: Callable[[Tree], T],
         *,
-        max_worker: Optional[int] = None,
+        max_worker: int | None = None,
         verbose: bool = False,
     ) -> Iterator[T]:
         """Map a function to all trees in the population.
 
-        This is a straightforward interface for parallelizing
-        computations. The parameters are intentionally kept simple and
-        user-friendly. For more advanced control, consider using
-        `concurrent.futures` directly.
+        This is a straightforward interface for parallelizing computations. The
+        parameters are intentionally kept simple and user-friendly. For more advanced
+        control, consider using `concurrent.futures` directly.
         """
 
         trees = (t for t in self.trees)
@@ -226,11 +218,11 @@ class Population:
         cls,
         root: str,
         ext: str = ".eswc",
-        extra_cols: Optional[Iterable[str]] = None,
+        extra_cols: Iterable[str] | None = None,
         **kwargs,
     ) -> Self:
         extra_cols = list(extra_cols) if extra_cols is not None else []
-        extra_cols.extend(k for k, t in eswc_cols)
+        extra_cols.extend(k for k, _ in eswc_cols)
         return cls.from_swc(root, ext, extra_cols=extra_cols, **kwargs)
 
     @staticmethod
@@ -253,23 +245,21 @@ class Populations:
     labels: list[str]
 
     def __init__(
-        self, populations: Iterable[Population], labels: Optional[Iterable[str]] = None
+        self, populations: Iterable[Population], labels: Iterable[str] | None = None
     ) -> None:
         self.len = min(len(p) for p in populations)
         self.populations = list(populations)
 
         labels = list(labels) if labels is not None else ["" for i in populations]
-        assert len(labels) == len(
-            self.populations
-        ), f"got {len( self.populations)} populations, but has {len(labels)} labels"
+        assert len(labels) == len(self.populations), (
+            f"got {len(self.populations)} populations, but has {len(labels)} labels"
+        )
         self.labels = labels
 
-    # fmt:off
     @overload
     def __getitem__(self, key: slice) -> list[list[Tree]]: ...
     @overload
     def __getitem__(self, key: int) -> list[Tree]: ...
-    # fmt:on
     def __getitem__(self, key):
         return [p[key] for p in self.populations]
 
@@ -299,24 +289,18 @@ class Populations:
         ext: str = ".swc",
         intersect: bool = True,
         check_same: bool = False,
-        labels: Optional[Iterable[str]] = None,
+        labels: Iterable[str] | None = None,
         **kwargs,
     ) -> Self:
         """Get population from dirs.
 
-        Parameters
-        ----------
-        roots : List of str
-        intersect : bool, default `True`
-            Take the intersection of these populations.
-        check_same : bool, default `False`
-            Check if the directories contains the same swc.
-        labels : List of str, optional
-            Label of populations.
-        **kwargs : Any
-            Forwarding to `Population`.
+        Args:
+            roots: List of str
+            intersect: Take the intersection of these populations.
+            check_same: Check if the directories contains the same swc.
+            labels: Label of populations.
+            **kwargs: Forwarding to `Population`.
         """
-
         fs = [Population.find_swcs(d, ext=ext, relpath=True) for d in roots]
         if intersect:
             inter = list(reduce(lambda a, b: set(a).intersection(set(b)), fs))
@@ -339,13 +323,13 @@ class Populations:
     def from_eswc(
         cls,
         roots: Iterable[str],
-        extra_cols: Optional[Iterable[str]] = None,
+        extra_cols: Iterable[str] | None = None,
         *,
         ext: str = ".eswc",
         **kwargs,
     ) -> Self:
         extra_cols = list(extra_cols) if extra_cols is not None else []
-        extra_cols.extend(k for k, t in eswc_cols)
+        extra_cols.extend(k for k, _ in eswc_cols)
         return cls.from_swc(roots, extra_cols=extra_cols, ext=ext, **kwargs)
 
 
@@ -360,11 +344,8 @@ def _get_idx(key: int, length: int) -> int:
 
 
 # experimental
-def filter_population(
-    pop: Population, predicate: Callable[[Tree], bool]
-) -> "Population":
+def filter_population(pop: Population, predicate: Callable[[Tree], bool]) -> Population:
     """Filter trees in the population."""
-
     # TODO: how to avoid load trees
     idx = [i for i, t in enumerate(pop) if predicate(t)]
     return Population(NestTrees(pop.trees, idx), root=pop.root)

swcgeom/core/swc.py

@@ -19,7 +19,7 @@ import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from copy import deepcopy
-from typing import Any, Optional, TypeVar, overload
+from typing import Any, TypeVar, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -145,17 +145,28 @@ class SWCLike(ABC):
         """Get the number of edges."""
         return self.number_of_nodes() - 1  # for tree structure: n = e + 1
 
-    # fmt: off
     @overload
-    def to_swc(self, fname: str, *, extra_cols: list[str] | None = ..., source: bool | str = ..., id_offset: int = ...) -> None: ...
+    def to_swc(
+        self,
+        fname: str,
+        *,
+        extra_cols: list[str] | None = ...,
+        source: bool | str = ...,
+        id_offset: int = ...,
+    ) -> None: ...
     @overload
-    def to_swc(self, *, extra_cols: list[str] | None = ..., source: bool | str = ..., id_offset: int = ...) -> str: ...
-    # fmt: on
     def to_swc(
         self,
-        fname: Optional[str] = None,
         *,
-        extra_cols: Optional[list[str]] = None,
+        extra_cols: list[str] | None = ...,
+        source: bool | str = ...,
+        id_offset: int = ...,
+    ) -> str: ...
+    def to_swc(
+        self,
+        fname: str | None = None,
+        *,
+        extra_cols: list[str] | None = None,
         source: bool | str = True,
         comments: bool = True,
         id_offset: int = 1,
@@ -183,17 +194,15 @@ class SWCLike(ABC):
 
         return None
 
-    # fmt: off
     @overload
     def to_eswc(self, fname: str, **kwargs) -> None: ...
     @overload
-    def to_eswc(self, **kwargs) -> str: ...
-    # fmt: on
+    def to_eswc(self, fname: None = ..., **kwargs) -> str: ...
     def to_eswc(
         self,
-        fname: Optional[str] = None,
-        swc_path: Optional[str] = None,
-        extra_cols: Optional[list[str]] = None,
+        fname: str | None = None,
+        swc_path: str | None = None,
+        extra_cols: list[str] | None = None,
         **kwargs,
     ) -> str | None:
         if swc_path is None:
@@ -205,7 +214,7 @@ class SWCLike(ABC):
             fname = swc_path
 
         extra_cols = extra_cols or []
-        extra_cols.extend(k for k, t in eswc_cols)
+        extra_cols.extend(k for k, _ in eswc_cols)
         return self.to_swc(fname, extra_cols=extra_cols, **kwargs)  # type: ignore
 
 
@@ -221,8 +230,8 @@ class DictSWC(SWCLike):
         self,
         *,
         source: str = "",
-        comments: Optional[Iterable[str]] = None,
-        names: Optional[SWCNames] = None,
+        comments: Iterable[str] | None = None,
+        names: SWCNames | None = None,
         **kwargs: npt.NDArray,
     ):
         super().__init__()

swcgeom/core/swc_utils/__init__.py

@@ -15,12 +15,10 @@
 
 """SWC format utils.
 
-Notes
------
-This module provides a bunch of methods to manipulating swc files, they
-are always trivial and unstabled, so we are NOT export it by default.
-If you use the method here, please review the code more frequently, we
-will try to flag all breaking changes but NO promises.
+NOTE: This module provides a bunch of methods to manipulating swc files, they are
+always trivial and unstabled, so we are NOT export it by default. If you use the method
+here, please review the code more frequently, we will try to flag all breaking changes
+but NO promises.
 """
 
 from swcgeom.core.swc_utils.assembler import *  # noqa: F403

swcgeom/core/swc_utils/assembler.py

@@ -15,10 +15,7 @@
 
 """Assemble lines to swc.
 
-Notes
------
-This module is deprecated, please use `~.transforms.LinesToTree`
-instead.
+NOTE: This module is deprecated, please use `~.transforms.LinesToTree` instead.
 """
 
 __all__ = ["assemble_lines", "try_assemble_lines"]
@@ -30,11 +27,9 @@ def assemble_lines(*args, **kwargs):
     .. deprecated:: 0.15.0
         Use :meth:`~.transforms.LinesToTree` instead.
     """
-
     raise DeprecationWarning(
-        "`assemble_lines` has been replaced by "
-        "`~.transforms.LinesToTree` because it can be easy assemble "
-        "with other tansformations.",
+        "`assemble_lines` has been replaced by `~.transforms.LinesToTree` because it "
+        "can be easy assemble with other tansformations.",
     )
 
 
@@ -44,9 +39,7 @@ def try_assemble_lines(*args, **kwargs):
     .. deprecated:: 0.15.0
         Use :meth:`~.transforms.LinesToTree` instead.
     """
-
     raise DeprecationWarning(
-        "`try_assemble_lines` has been replaced by "
-        "`~.transforms.LinesToTree` because it can be easy assemble "
-        "with other tansformations.",
+        "`try_assemble_lines` has been replaced by `~.transforms.LinesToTree` because "
+        "it can be easy assemble with other tansformations.",
     )

swcgeom/core/swc_utils/base.py

@@ -16,7 +16,7 @@
 """Base SWC format utils."""
 
 from collections.abc import Callable
-from typing import Literal, NamedTuple, Optional, TypeVar, overload
+from typing import Literal, NamedTuple, TypeVar, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -34,7 +34,8 @@ __all__ = [
     "traverse",
 ]
 
-T, K = TypeVar("T"), TypeVar("K")
+T = TypeVar("T")
+K = TypeVar("K")
 Topology = tuple[npt.NDArray[np.int32], npt.NDArray[np.int32]]  # (id, pid)
 
 
@@ -56,17 +57,15 @@ class SWCNames(NamedTuple):
 swc_names = SWCNames()
 
 
-def get_names(names: Optional[SWCNames] = None) -> SWCNames:
+def get_names(names: SWCNames | None = None) -> SWCNames:
     return names or swc_names
 
 
 class SWCTypes(NamedTuple):
     """SWC format types.
 
-    See Also
-    ---------
-    NeuroMoprho.org - What is SWC format?
-        https://neuromorpho.org/myfaq.jsp
+    See Also:
+        NeuroMoprho.org - What is SWC format?: https://neuromorpho.org/myfaq.jsp
     """
 
     undefined: int = 0
@@ -82,17 +81,17 @@ class SWCTypes(NamedTuple):
 swc_types = SWCTypes()
 
 
-def get_types(types: Optional[SWCTypes] = None) -> SWCTypes:
+def get_types(types: SWCTypes | None = None) -> SWCTypes:
     return types or swc_types
 
 
-def get_topology(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> Topology:
+def get_topology(df: pd.DataFrame, *, names: SWCNames | None = None) -> Topology:
     names = get_names(names)
     return (df[names.id].to_numpy(), df[names.pid].to_numpy())
 
 
 def get_dsu(
-    df: pd.DataFrame, *, names: Optional[SWCNames] = None
+    df: pd.DataFrame, *, names: SWCNames | None = None
 ) -> npt.NDArray[np.int32]:
     """Get disjoint set union."""
     names = get_names(names)
@@ -116,36 +115,46 @@ def get_dsu(
     return dsu
 
 
-# fmt: off
 @overload
-def traverse(topology: Topology, *, enter: Callable[[int, T | None], T], root: int | np.integer = ..., mode: Literal["dfs"] = ...) -> None: ...
+def traverse(
+    topology: Topology,
+    *,
+    enter: Callable[[int, T | None], T],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
+) -> None: ...
 @overload
-def traverse(topology: Topology, *, leave: Callable[[int, list[K]], K], root: int | np.integer = ..., mode: Literal["dfs"] = ...) -> K: ...
+def traverse(
+    topology: Topology,
+    *,
+    leave: Callable[[int, list[K]], K],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
+) -> K: ...
 @overload
 def traverse(
-    topology: Topology, *, enter: Callable[[int, T | None], T], leave: Callable[[int, list[K]], K],
-    root: int | np.integer = ..., mode: Literal["dfs"] = ...,
+    topology: Topology,
+    *,
+    enter: Callable[[int, T | None], T],
+    leave: Callable[[int, list[K]], K],
+    root: int | np.integer = ...,
+    mode: Literal["dfs"] = ...,
 ) -> K: ...
-# fmt: on
 def traverse(topology: Topology, *, mode="dfs", **kwargs):
     """Traverse nodes.
 
-    Parameters
-    ----------
-    enter : (id: int, parent: T | None) => T, optional
-        The callback when entering node, which accepts two parameters,
-        the current node id and the return value of it parent node. In
-        particular, the root node receives an `None`.
-    leave : (id: int, children: list[T]) => T, optional
-        The callback when leaving node. When leaving a node, subtree
-        has already been traversed. Callback accepts two parameters,
-        the current node id and list of the return value of children,
-        In particular, the leaf node receives an empty list.
-    root : int, default to `0`
-        Start from the root node of the subtree
-    mode : `dfs`, default to `dfs`
+    Args:
+    enter: (id: int, parent: T | None) => T, optional
+        The callback when entering node, which accepts two parameters, the current node
+        id and the return value of it parent node. In particular, the root node
+        receives an `None`.
+    leave: (id: int, children: list[T]) => T, optional
+        The callback when leaving node. When leaving a node, subtree has already been
+        traversed. Callback accepts two parameters, the current node id and list of the
+        return value of children, In particular, the leaf node receives an empty list.
+    root: Start from the root node of the subtree
+    mode: The traverse mode, only support "dfs" now.
     """
-
     match mode:
         case "dfs":
             return _traverse_dfs(topology, **kwargs)

swcgeom/core/swc_utils/checker.py

@@ -16,7 +16,6 @@
 """Check common"""
 
 from collections import defaultdict
-from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -36,14 +35,13 @@ __all__ = [
 ]
 
 
-def is_single_root(df: pd.DataFrame, *, names: Optional[SWCNames] = None) -> bool:
+def is_single_root(df: pd.DataFrame, *, names: SWCNames | None = None) -> bool:
     """Check is it only one root."""
     return len(np.unique(get_dsu(df, names=names))) == 1
 
 
 def is_bifurcate(topology: Topology, *, exclude_root: bool = True) -> bool:
     """Check is it a bifurcate topology."""
-
     children = defaultdict(list)
     for idx, pid in zip(*topology):
         children[pid].append(idx)
@@ -59,8 +57,7 @@ def is_bifurcate(topology: Topology, *, exclude_root: bool = True) -> bool:
 def is_sorted(topology: Topology) -> bool:
     """Check is it sorted.
 
-    In a sorted topology, parent samples should appear before any child
-    samples.
+    In a sorted topology, parent samples should appear before any child samples.
     """
     flag = True
 
@@ -103,20 +100,18 @@ def check_single_root(*args, **kwargs) -> bool:
     .. deprecated:: 0.5.0
         Use :meth:`is_single_root` instead.
     """
-
     return is_single_root(*args, **kwargs)
 
 
 @deprecated("Use `is_bifurcate` instead")
 def is_binary_tree(
-    df: pd.DataFrame, exclude_root: bool = True, *, names: Optional[SWCNames] = None
+    df: pd.DataFrame, exclude_root: bool = True, *, names: SWCNames | None = None
 ) -> bool:
     """Check is it a binary tree.
 
     .. deprecated:: 0.8.0
         Use :meth:`is_bifurcate` instead.
     """
-
     names = get_names(names)
     topo = (df[names.id].to_numpy(), df[names.pid].to_numpy())
     return is_bifurcate(topo, exclude_root=exclude_root)

swcgeom/core/swc_utils/io.py

@@ -18,7 +18,7 @@
 import re
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Literal, Optional
+from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -39,41 +39,33 @@ __all__ = ["read_swc", "to_swc"]
 
 def read_swc(
     swc_file: PathOrIO,
-    extra_cols: Optional[Iterable[str]] = None,
+    extra_cols: Iterable[str] | None = None,
     fix_roots: Literal["somas", "nearest", False] = False,
     sort_nodes: bool = False,
     reset_index: bool = True,
     *,
     encoding: Literal["detect"] | str = "utf-8",
-    names: Optional[SWCNames] = None,
+    names: SWCNames | None = None,
 ) -> tuple[pd.DataFrame, list[str]]:
     """Read swc file.
 
-    Parameters
-    ----------
-    swc_file : PathOrIO
-        Path of swc file, the id should be consecutively incremented.
-    extra_cols : Iterable[str], optional
-        Read more cols in swc file.
-    fix_roots : `somas`|`nearest`|False, default `False`
-        Fix multiple roots.
-    sort_nodes : bool, default `False`
-        Sort the indices of neuron tree, the index for parent are
-        always less than children.
-    reset_index : bool, default `True`
-        Reset node index to start with zero, DO NOT set to false if
-        you are not sure what will happend.
-    encoding : str | 'detect', default `utf-8`
-        The name of the encoding used to decode the file. If is
-        `detect`, we will try to detect the character encoding.
-    names : SWCNames, optional
-
-    Returns
-    -------
-    df : ~pandas.DataFrame
-    comments : List of string
+    NOTE: the id should be consecutively incremented.
+
+    Args:
+        extra_cols: Read more cols in swc file.
+        fix_roots: Fix multiple roots.
+        sort_nodes: Sort the indices of neuron tree.
+            After sorting the nodes, the index for each parent are always less than
+            that of its children.
+        reset_index: Reset node index to start with zero.
+            DO NOT set to false if you are not sure what will happened.
+        encoding: The name of the encoding used to decode the file.
+            If is `detect`, we will try to detect the character encoding.
+
+    Returns:
+        df: ~pandas.DataFrame
+        comments: List of string
     """
-
     names = get_names(names)
     df, comments = parse_swc(
         swc_file, names=names, extra_cols=extra_cols, encoding=encoding
@@ -110,10 +102,10 @@ def read_swc(
 def to_swc(
     get_ndata: Callable[[str], npt.NDArray],
     *,
-    extra_cols: Optional[Iterable[str]] = None,
+    extra_cols: Iterable[str] | None = None,
     id_offset: int = 1,
-    comments: Optional[Iterable[str]] = None,
-    names: Optional[SWCNames] = None,
+    comments: Iterable[str] | None = None,
+    names: SWCNames | None = None,
 ) -> Iterable[str]:
     """Convert to swc format."""
 
@@ -156,21 +148,14 @@ def parse_swc(
 ) -> tuple[pd.DataFrame, list[str]]:
     """Parse swc file.
 
-    Parameters
-    ----------
-    fname : PathOrIO
-    names : SWCNames
-    extra_cols : List of str, optional
-    encoding : str | 'detect', default `utf-8`
-        The name of the encoding used to decode the file. If is
-        `detect`, we will try to detect the character encoding.
-
-    Returns
-    -------
-    df : ~pandas.DataFrame
-    comments : List of string
-    """
+    Args:
+        encoding: The name of the encoding used to decode the file.
+            If is `detect`, we will try to detect the character encoding.
 
+    Returns:
+        df: ~pandas.DataFrame
+        comments: List of string
+    """
     # pylint: disable=too-many-locals
     extras = list(extra_cols) if extra_cols else []
 
@@ -208,7 +193,7 @@ def parse_swc(
                 if (match := re_swc.search(line)) is not None:
                     if flag and match.group(last_group):
                         warnings.warn(
-                            f"some fields are ignored in row {i+1} of `{fname}`"
+                            f"some fields are ignored in row {i + 1} of `{fname}`"
                         )
                         flag = False
 
@@ -219,10 +204,10 @@ def parse_swc(
                     if not comment.startswith(ignored_comment):
                         comments.append(comment)
                 elif not line.isspace():
-                    raise ValueError(f"invalid row {i+1} in `{fname}`")
+                    raise ValueError(f"invalid row {i + 1} in `{fname}`")
         except UnicodeDecodeError as e:
             raise ValueError(
-                f"decode failed, try to enable auto detect `encoding='detect'`"
+                "decode failed, try to enable auto detect `encoding='detect'`"
             ) from e
 
     df = pd.DataFrame.from_dict(dict(zip(keys, vals)))