From 0533ad53653167e6d8b20335fc419a7a202e63ed Mon Sep 17 00:00:00 2001
From: Joe Zuntz <joezuntz@googlemail.com>
Date: Thu, 28 Sep 2023 11:03:15 +0100
Subject: [PATCH] Add command for nicer docstring

---
 ceci/stage.py       | 30 ++++++++++++++++++++++++++++++
 tests/test_stage.py |  2 ++
 2 files changed, 32 insertions(+)

diff --git a/ceci/stage.py b/ceci/stage.py
index 1915a93..f146876 100644
--- a/ceci/stage.py
+++ b/ceci/stage.py
@@ -333,6 +333,16 @@ def __init_subclass__(cls, **kwargs):
         # Find the absolute path to the class defining the file
         path = pathlib.Path(filename).resolve()
 
+        # Add a description of the parameters to the end of the docstring
+        if stage_is_complete:
+            config_text = cls._describe_configuration_text()
+            if cls.__doc__ is None:
+                cls.__doc__ = f"Stage {cls.name}\n\nConfiguration Parameters:\n{config_text}"
+            else:
+                # strip any existing configuration text from parent classes that is at the end of the doctring
+                cls.__doc__ = cls.__doc__.split("Configuration Parameters:")[0]
+                cls.__doc__ += f"\n\nConfiguration Parameters:\n{config_text}"
+
         # Register the class
         if stage_is_complete:
             cls.pipeline_stages[cls.name] = (cls, path)
@@ -398,6 +408,26 @@ def get_module(cls):
         """
         return cls.pipeline_stages[cls.name][0].__module__
 
+    @classmethod
+    def describe_configuration(cls):
+        print(cls._describe_configuration_text())
+
+    @classmethod
+    def _describe_configuration_text(cls):
+        s = []
+        for name, val in cls.config_options.items():
+            if isinstance(val, StageParameter):
+                if val.required:
+                    txt = f"[{val.dtype.__name__}]: {val._help}  (required)"
+                else:
+                    txt = f"[{val.dtype.__name__}]: {val._help} (default={val.default})"
+            elif isinstance(val, type):
+                txt = f"[{val.__name__}]: (required)"
+            else:
+                txt = f"[{type(val).__name__}]: (default={val})"
+            s.append(f"{name} {txt} ")
+        return '\n'.join(s)
+
     @classmethod
     def usage(cls):  # pragma: no cover
         """
diff --git a/tests/test_stage.py b/tests/test_stage.py
index eec565c..e9bc705 100644
--- a/tests/test_stage.py
+++ b/tests/test_stage.py
@@ -145,6 +145,8 @@ def run(self):
     assert stage_1.config.b == 'puffins are not extinct?'
     assert 1 in stage_1.config.c
     assert len(stage_1.config.d) == 0
+    assert "c [list]: a list (default=[1, 2, 3])" in stage_1.__doc__
+    stage_1.describe_configuration()
     
     cmd = "TestStage", "--a", "6", "--b", "puffins are not extinct?", "--inp", "dummy"
     stage_1_cmd = TestStage(TestStage.parse_command_line(cmd))