add runtime profiling as option to the command line client

Author: martin-boettcher

docs/cli.md

@@ -22,6 +22,7 @@ Options:
                          'target_dir' configuration field.
   --dry-run              Run the tool without creating, changing, or deleting
                          any files.
+  --profiling TEXT       Path for profiling output. Switches on profiling.  
   --help-config json|md  Show configuration help and exit.
   --help                 Show this message and exit.
 ```

tests/test_cli.py

@@ -1,7 +1,7 @@
 # Copyright © 2024 Norman Fomferra
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
-
+import os
 import unittest
 from click.testing import CliRunner
 
@@ -104,3 +104,27 @@ def test_some_slices_and_no_target(self):
             "Error: Missing 'target_dir' in configuration\n", result.output
         )
         self.assertFalse(FileObj("memory://target.zarr").exists())
+
+    def test_profiling(self):
+        make_test_dataset(uri="memory://slice-1.zarr")
+        make_test_dataset(uri="memory://slice-2.zarr")
+        make_test_dataset(uri="memory://slice-3.zarr")
+
+        runner = CliRunner()
+        # noinspection PyTypeChecker
+        result = runner.invoke(
+            zappend,
+            [
+                "--profiling",
+                "prof.out",
+                "--target",
+                "memory://target.zarr",
+                "memory://slice-1.zarr",
+                "memory://slice-2.zarr",
+                "memory://slice-3.zarr",
+            ],
+        )
+        self.assertEqual("", result.output)
+        self.assertEqual(0, result.exit_code)
+        self.assertTrue(FileObj("prof.out").exists())
+        os.unlink("prof.out")

zappend/api.py

@@ -1,8 +1,10 @@
 # Copyright © 2024 Norman Fomferra
 # Permissions are hereby granted under the terms of the MIT License:
 # https://opensource.org/licenses/MIT.
-
-from typing import Iterable, Any
+from cProfile import Profile
+from io import StringIO
+from pstats import Stats
+from typing import Iterable, Any, Literal
 
 from .config import ConfigItem
 from .config import ConfigLike
@@ -15,6 +17,7 @@
 from .slice import SliceSource
 from .slice import to_slice_factories
 from .slice import to_slice_factory
+from .log import logger
 
 
 __all__ = [
@@ -31,9 +34,14 @@
     "to_slice_factory",
 ]
 
+_TOTAL_TIME: Literal["tottime"] = "tottime"
+
 
 def zappend(
-    slices: Iterable[SliceObj | SliceFactory], config: ConfigLike = None, **kwargs: Any
+    slices: Iterable[SliceObj | SliceFactory],
+    config: ConfigLike = None,
+    profiling=None,
+    **kwargs: Any,
 ):
     """
     Robustly create or update a Zarr dataset from dataset slices.
@@ -57,8 +65,25 @@ def zappend(
             May be a file path or URI, a ``dict``, ``None``, or a sequence of
             the aforementioned. If a sequence is used, subsequent configurations
             are incremental to the previous ones.
+        profiling: Path for profiling output.
+            Switches on profiling
         kwargs: Additional configuration parameters.
             Can be used to pass or override configuration values in *config*.
     """
     processor = Processor(config, **kwargs)
+
+    if profiling:
+        profile = Profile()
+        logger.info(f"profiling ...")
+        profile.enable()
+
     processor.process_slices(slices)
+
+    if profiling:
+        profile.disable()
+        results = StringIO()
+        stats = Stats(profile, stream=results)
+        stats.sort_stats(_TOTAL_TIME).print_stats()
+        with open(profiling, "w") as f:
+            f.write(results.getvalue())
+        logger.info(f"profiling output written to {profiling}")

zappend/cli.py

@@ -28,6 +28,7 @@
     is_flag=True,
     help="Run the tool without creating, changing," " or deleting any files.",
 )
+@click.option("--profiling", help="Path for profiling output. Switches on profiling.")
 @click.option(
     "--help-config",
     metavar="json|md",
@@ -39,6 +40,7 @@ def zappend(
     config: tuple[str, ...],
     target: str | None,
     dry_run: bool,
+    profiling: str,
     help_config: str | None,
 ):
     """Create or update a Zarr datacube TARGET from slice datasets SLICES.
@@ -61,7 +63,13 @@ def zappend(
 
     # noinspection PyBroadException
     try:
-        zappend(slices, config=config, target_dir=target, dry_run=dry_run)
+        zappend(
+            slices,
+            config=config,
+            target_dir=target,
+            dry_run=dry_run,
+            profiling=profiling,
+        )
     except BaseException as e:
         raise click.ClickException(f"{e}") from e