From 76b66be9519785f5c38d8931186ede86174e79a7 Mon Sep 17 00:00:00 2001
From: Ziwen Liu <ziwen.liu@czbiohub.org>
Date: Thu, 10 Oct 2024 10:27:59 -0700
Subject: [PATCH] improve type hint and docstring

---
 viscy/cli.py                  | 7 ++++---
 viscy/translation/__main__.py | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/viscy/cli.py b/viscy/cli.py
index ec264082..3de758e0 100644
--- a/viscy/cli.py
+++ b/viscy/cli.py
@@ -22,7 +22,7 @@ def subcommands() -> dict[str, set[str]]:
         subcommands["export"] = {"model", "dataloaders", "datamodule"}
         return subcommands
 
-    def add_arguments_to_parser(self, parser):
+    def add_arguments_to_parser(self, parser) -> None:
         parser.set_defaults(
             {
                 "trainer.logger": lazy_instance(
@@ -40,16 +40,17 @@ def run_cli(
     model_class: type[LightningModule],
     datamodule_class: type[LightningDataModule],
     trainer_class: type[VisCyTrainer],
-):
+) -> None:
     """
     Main Lightning CLI entry point.
+    Parse log level and set TF32 precision.
 
     Parameters
     ----------
     cli_class : type[LightningCLI]
         Lightning CLI class
     model_class : type[LightningModule]
-        Lightning module class
+        Lightning module class. Not used in preprocess mode.
     datamodule_class : type[LightningDataModule]
         Lightning datamodule class
     trainer_class : type[VisCyTrainer]
diff --git a/viscy/translation/__main__.py b/viscy/translation/__main__.py
index a360fbdf..597e48d4 100644
--- a/viscy/translation/__main__.py
+++ b/viscy/translation/__main__.py
@@ -9,7 +9,7 @@
 class TranslationCLI(VisCyCLI):
     """Extending lightning CLI arguments and defualts."""
 
-    def add_arguments_to_parser(self, parser):
+    def add_arguments_to_parser(self, parser) -> None:
         super().add_arguments_to_parser(parser)
         if "preprocess" not in sys.argv:
             parser.link_arguments("data.yx_patch_size", "model.example_input_yx_shape")