Merge pull request #1152 from wright-group/wt-tree

reworking wright-tools entry points

CHANGELOG.md

@@ -5,6 +5,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ## [Unreleased]
 
+### Added
+- `wt5` entry point: methods `load`, `explore`, `glob`, `tree`
+- `wt-convert` entry point: convert a number to different units
+- new glob utilities in `kit`
+
+### Changed
+- `wt-tree` entry point removed.  Use `wt5 tree` instead.
+
 ## [3.5.1]
 
 ### Added

WrightTools/cli/__init__.py

@@ -0,0 +1,2 @@
+from . import _units
+from . import _wt5

WrightTools/cli/_units.py

@@ -0,0 +1,50 @@
+# --- import --------------------------------------------------------------------------------------
+
+
+import click
+from WrightTools import __version__ as __wt_version__
+from WrightTools.units import is_valid_conversion, get_valid_conversions, convert
+from WrightTools.exceptions import UnitsError
+
+
+# --- define --------------------------------------------------------------------------------------
+
+
+@click.command(name="convert", help="convert numbers to different units.")
+@click.version_option(__wt_version__, prog_name="WrightTools")
+@click.argument("number", type=float, nargs=1)
+@click.argument("unit", nargs=1)
+@click.argument("destination_unit", default=None, nargs=-1)
+def cli(number, unit, destination_unit=None):
+    """Convert numbers to different units."""
+
+    if int(number) == number:
+        number = int(number)
+    sig_figs = len(str(number))
+    sig_figs -= 1 if "." in str(number) else 0
+
+    def fmt(new):
+        exponent = int(f"{new:e}".split("e")[1])
+        if exponent > 6 or exponent < -3:
+            return f"{new:{sig_figs}e}"
+        else:  # if a "normal" size number
+            if sig_figs - exponent <= 0:
+                return f"{int(round(new, sig_figs-exponent))}"
+            else:
+                return f"{round(new, sig_figs-exponent)}"
+
+    if len(destination_unit):  # units provided
+        destination_unit = destination_unit[0]
+        if not is_valid_conversion(unit, destination_unit):
+            raise UnitsError(get_valid_conversions(unit), destination_unit)
+        new = convert(number, unit, destination_unit)
+        print(f"{number} {unit} = {fmt(new)} {destination_unit}")
+    else:
+        valid_units = get_valid_conversions(unit)
+        for d_unit in valid_units:
+            new = convert(number, unit, d_unit)
+            print(f"{fmt(new)} {d_unit}")
+
+
+if __name__ == "__main__":
+    cli()

WrightTools/cli/_wt5.py

@@ -0,0 +1,150 @@
+# --- import --------------------------------------------------------------------------------------
+
+
+import click
+import WrightTools as wt
+
+
+# --- define --------------------------------------------------------------------------------------
+
+
+@click.group()
+@click.version_option(wt.__version__, prog_name="WrightTools")
+def cli():
+    pass
+
+
+@cli.command(name="tree", help="Print a given data tree.")
+@click.argument("path", nargs=1)
+@click.option(
+    "--internal_path", default="/", help="specify a path internal to the file.  Defaults to root"
+)
+@click.option("--depth", "-d", "-L", type=int, default=9, help="Depth to print.")
+@click.option("--verbose", "-v", is_flag=True, default=False, help="Print a more detailed tree.")
+def tree(path, internal_path, depth=9, verbose=False):
+    # open the object
+    obj = wt.open(path)[internal_path]
+
+    if isinstance(obj, wt.Data):
+        obj.print_tree(verbose=verbose)
+    else:
+        obj.print_tree(verbose=verbose, depth=depth)
+
+
+@cli.command(name="load", help="Open a python cli with the object pre-loaded.")
+@click.argument("path")
+def load(path):
+    import code
+
+    shell = code.InteractiveConsole()
+    _interact(shell, path)
+
+
+@cli.command(name="glob", help="Find paths to all wt5 files within a directory.")
+@click.option(
+    "--directory", "-d", default=None, help="Directory to scan.  Defaults to current directory."
+)
+@click.option("--no-recursion", is_flag=True, help="Turns recursive scan off.")
+def glob(directory=None, no_recursion=False):
+    for path in wt.kit.glob_wt5s(directory, not no_recursion):
+        print(str(path))
+
+
+@cli.command(name="explore", help="Scan a directory and survey the wt5 objects found.")
+@click.option(
+    "--directory", "-d", default=None, help="Directory to scan.  Defaults to current directory."
+)
+@click.option("--no-recursion", is_flag=True, help="Turns recursive scan off.")
+# TODO: formatting options (e.g. json)?
+def scan(directory=None, no_recursion=False):
+    import os, code
+    from rich.live import Live
+    from rich.table import Table
+
+    if directory is None:
+        directory = os.getcwd()
+
+    table = Table(title=directory)
+    table.add_column("", justify="right")  # index
+    table.add_column("path", max_width=60, no_wrap=True, style="blue")
+    table.add_column("size (MB)", justify="center", style="blue")
+    table.add_column("created", max_width=30, style="blue")
+    table.add_column("name", style="blue")
+    table.add_column("shape", style="blue")
+    table.add_column("axes", max_width=50, style="blue")
+    table.add_column("variables", style="blue")
+    table.add_column("channels", style="blue")
+
+    update_title = lambda n: directory + f" ({n} wt5 file{'s' if n != 1 else None} found)"
+    paths = []
+
+    with Live(table) as live:
+        for i, path in enumerate(wt.kit.glob_wt5s(directory, not no_recursion)):
+            desc = wt.kit.describe_wt5(path)
+            desc["filesize"] = f"{os.path.getsize(path) / 1e6:.1f}"
+            path = path.relative_to(directory)
+            paths.append(path)
+            desc["path"] = (
+                f"[link={path.parent}]{path.parent}[/link]" + r"\\"
+                if str(path.parent) != "."
+                else ""
+            )
+            desc["path"] += f"[bold]{path.name}[/bold]"
+            # desc["path"] = f"[link={str(path)}]{path}[/link]"
+            row = [f"{i}"] + [
+                str(desc[k])
+                for k in ["path", "filesize", "created", "name", "shape", "axes", "nvars", "nchan"]
+            ]
+            table.title = update_title(i + 1)
+            table.add_row(*row)
+            live.update(table)
+
+    # give option to interact
+    def raise_sys_exit():
+        raise SystemExit
+
+    shell = code.InteractiveConsole(locals={"exit": raise_sys_exit, "quit": raise_sys_exit})
+
+    while True:
+        msg = shell.raw_input(
+            " ".join(
+                [
+                    "Specify an index to load that entry.",
+                    "Use `t` to rerender table.",
+                    "Use no argument to exit.",
+                ]
+            )
+        )
+        if msg == "t":
+            with Live(table) as live:
+                pass
+            continue
+        try:
+            valid = 0 <= int(msg) < len(paths)
+        except ValueError:
+            break
+        if valid:
+            print("interacting...")
+            _interact(shell, str(paths[int(msg)]))
+            continue
+
+
+def _interact(shell, path):
+    lines = [
+        "import WrightTools as wt",
+        "import matplotlib.pyplot as plt",
+        f"d = wt.open(r'{path}')",
+    ]
+
+    [shell.push(line) for line in lines]
+    banner = "--- INTERACTING --- (to continue, call exit() or quit())\n"
+    banner += "\n".join([">>> " + line for line in lines])
+
+    try:
+        shell.interact(banner=banner)
+    except SystemExit:
+        pass
+
+
+if __name__ == "__main__":
+    cli()

WrightTools/data/_data.py

@@ -272,30 +272,30 @@ def print_leaves(prefix, lis, vline=True):
 
         if verbose:
             # axes
-            print(prefix + "├── axes")
+            print(prefix + f"├── axes ({len(self.axes)})")
             print_leaves(prefix, self.axes)
             # constants
-            print(prefix + "├── constants")
+            print(prefix + f"├── constants ({len(self.constants)})")
             print_leaves(prefix, self.constants)
             # variables
-            print(prefix + "├── variables")
+            print(prefix + f"├── variables ({len(self.variables)})")
             print_leaves(prefix, self.variables)
             # channels
-            print(prefix + "└── channels")
+            print(prefix + f"└── channels ({len(self.channels)})")
             print_leaves(prefix, self.channels, vline=False)
         else:
             # axes
-            s = "axes: "
+            s = f"axes ({len(self.axes)}): "
             s += ", ".join(["{0} ({1})".format(a.expression, a.units) for a in self.axes])
             print(prefix + "├── " + s)
             # constants
-            s = "constants: "
+            s = f"constants ({len(self.constants)}): "
             s += ", ".join(
                 ["{0} ({1} {2})".format(a.expression, a.value, a.units) for a in self.constants]
             )
             print(prefix + "├── " + s)
             # channels
-            s = "channels: "
+            s = f"channels ({len(self.channels)}): "
             s += ", ".join(self.channel_names)
             print(prefix + "└── " + s)
 
@@ -1506,7 +1506,7 @@ def offset(
 
     def print_tree(self, *, verbose=True):
         """Print a ascii-formatted tree representation of the data contents."""
-        print("{0} ({1})".format(self.natural_name, self.filepath))
+        print("{0} ({1}) {2}".format(self.natural_name, self.filepath, self.shape))
         self._print_branch("", depth=0, verbose=verbose)
 
     def prune(self, keep_channels=True, *, verbose=True):

WrightTools/kit/__init__.py

@@ -14,3 +14,4 @@
 from ._timestamp import *
 from ._unicode import *
 from ._utilities import *
+from ._glob import *