feat(runtime): add display name, description, category and output nod…

…e to node docstring
Chaoses-Ib · Feb 7, 2024 · 6bf8647 · 6bf8647
1 parent c1be366
commit 6bf8647
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 15 deletions.
diff --git a/src/comfy_script/runtime/factory.py b/src/comfy_script/runtime/factory.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 from enum import Enum
+import textwrap
 from typing import Any
 
 from .. import astutil
@@ -417,21 +418,26 @@ def type_and_hint(type_info: str | list[str | bool], name: str = None, optional:
 
         # Docstring
         # TODO: Return names
-        # TODO: Display name
-        # TODO: Category
         # TODO: Min, max
         # TODO: Round?
-        # TODO: Indent
+        # TODO: Package (real mode)
         doc = (
-f'''    def {class_id}(
-        {f",{chr(10)}        ".join(inputs)}
-    ){output_type_hint}''')
+f'''```
+def {class_id}(
+    {f",{chr(10)}    ".join(inputs)}
+){output_type_hint}
+```''')
+        if info['display_name'] != class_id:
+            doc += f"\n{info['display_name']}\n"
 
-        quote = "'''" if "'''" not in doc else '"""'
-        c += (
-f"""    {quote}```
-{doc}
-    ```""")
+        # Almost no node has description
+        if info['description'] != '':
+            doc += f"\n{info['description']}\n"
+
+        doc += f"\nCategory: {info['category']}\n"
+
+        if info['output_node']:
+            doc += '\nOutput node.\n'
 
         for i, input in reversed(list(enumerate(inputs))):
             if '=' not in input:
@@ -441,10 +447,11 @@ def type_and_hint(type_info: str | list[str | bool], name: str = None, optional:
                         inputs[j] = input[:input.index('=')].rstrip() + ' | None'
                         removed = True
                 if removed:
-                    c += '\n    Use `None` to use default values of arguments that appear before any non-default argument.\n    '
+                    doc += '\nUse `None` to use default values of arguments that appear before any non-default argument.\n'
                 break
-
-        c += f'{quote}\n'
+
+        quote = "'''" if "'''" not in doc else '"""'
+        c += f"""{textwrap.indent(f'{quote}{doc}{quote}', '    ')}\n"""
 
         # __new__
         c += f'    def __new__(cls, {", ".join(inputs)}){output_type_hint}: ...\n'

diff --git a/tests/test_astutil.py b/tests/test_astutil.py
@@ -1,3 +1,4 @@
+'''hatch env run -e test pytest tests/test_astutil.py'''
 import pytest
 
 import comfy_script.astutil as astutil
@@ -9,4 +10,13 @@
     ("a'", "'''a\\''''"),
 ])
 def test_to_str(s, c):
-    assert astutil.to_str(s) == c
+    assert astutil.to_str(s) == c
+
+@pytest.mark.parametrize('fullname, module', [
+    ('PIL.Image.Image', 'PIL.Image'),
+    ('PIL.Image.Image.TEST', 'PIL.Image'),
+    ('PIL.Image', 'PIL'),
+    ('THONPY.Image', None),
+])
+def test_resolve_fullname(fullname, module):
+    assert getattr(astutil.find_spec_from_fullname(fullname), 'name', None) == module