BUG: Fixed import issues with Python stub interface files.

Previously, some wrapped classes inherited from classes defined in a seperate file. These classes were not imported. This fix allows the classes to be imported without requiring each file to know the name of the file that the class is located in. All proxies are imported in the _proxies.pyi file which can then be imported by different template classes. The Proxy and Template classes were split to allow the leading underscore to be removed from the proxies without causing extra 'fake' classes to appear for the user.
InsightSoftwareConsortium · Jan 31, 2022 · 9f6a17f · 9f6a17f
1 parent cff5823
commit 9f6a17f
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 19 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -637,10 +637,17 @@ set(ITK_DIR "${ITK_BINARY_DIR}")
 if(ITK_WRAPPING)
   add_subdirectory(Wrapping)
   if(ITK_WRAP_PYTHON)
-    set(ITK_STUB_IMPORTS "$ENV{ITK_STUB_IMPORTS}")
-    list(SORT ITK_STUB_IMPORTS CASE INSENSITIVE)
-    list(JOIN ITK_STUB_IMPORTS "\n" ITK_STUB_IMPORTS)
+    # Generate __init__.pyi file with wrapped templates/interfaces
+    set(ITK_STUB_TEMPLATE_IMPORTS "$ENV{ITK_STUB_TEMPLATE_IMPORTS}")
+    list(SORT ITK_STUB_TEMPLATE_IMPORTS CASE INSENSITIVE)
+    list(JOIN ITK_STUB_TEMPLATE_IMPORTS "\n" ITK_STUB_TEMPLATE_IMPORTS)
     configure_file(Wrapping/Generators/Python/itk/__init__.pyi.in Wrapping/Generators/Python/itk-stubs/__init__.pyi)
+
+    # Generate _proxies.pyi file with wrapped proxy classes and method
+    set(ITK_STUB_PROXY_IMPORTS "$ENV{ITK_STUB_PROXY_IMPORTS}")
+    list(SORT ITK_STUB_PROXY_IMPORTS CASE INSENSITIVE)
+    list(JOIN ITK_STUB_PROXY_IMPORTS "\n" ITK_STUB_PROXY_IMPORTS)
+    configure_file(Wrapping/Generators/Python/itk/_proxies.pyi.in Wrapping/Generators/Python/itk-stubs/_proxies.pyi)
   endif()
 endif()
 

diff --git a/Wrapping/Generators/Python/itk/__init__.pyi.in b/Wrapping/Generators/Python/itk/__init__.pyi.in
@@ -8,4 +8,4 @@ from .support.extras import *
 from .support.types import *
 
 # Begin stub imports
-${ITK_STUB_IMPORTS}
+${ITK_STUB_TEMPLATE_IMPORTS}
diff --git a/Wrapping/Generators/Python/itk/_proxies.pyi.in b/Wrapping/Generators/Python/itk/_proxies.pyi.in
@@ -0,0 +1,5 @@
+# Python proxy import file
+# This file allows the wrapped classes (both templates and proxies)
+# to import proxy classes without exposing the Proxy classes to a user
+# This also allows the imports to be done in other files without knowing the classes file name
+${ITK_STUB_PROXY_IMPORTS}
diff --git a/Wrapping/Generators/SwigInterface/CMakeLists.txt b/Wrapping/Generators/SwigInterface/CMakeLists.txt
@@ -333,8 +333,10 @@ macro(itk_end_wrap_module_swig_interface)
     list(APPEND typedef_in_files "${WRAPPER_LIBRARY_OUTPUT_DIR}/${module}SwigInterface.h.in")
     list(APPEND typedef_files "${WRAPPER_MASTER_INDEX_OUTPUT_DIR}/${module}SwigInterface.h")
     if(${module_prefix}_WRAP_PYTHON)
-      list(APPEND ITK_STUB_PYI_FILES "${ITK_STUB_DIR}/${module}.pyi")
-      set(ENV{ITK_STUB_IMPORTS} "$ENV{ITK_STUB_IMPORTS}from .${module} import *;")
+      list(APPEND ITK_STUB_PYI_FILES "${ITK_STUB_DIR}/${module}Template.pyi")
+      list(APPEND ITK_STUB_PYI_FILES "${ITK_STUB_DIR}/${module}Proxy.pyi")
+      set(ENV{ITK_STUB_TEMPLATE_IMPORTS} "$ENV{ITK_STUB_TEMPLATE_IMPORTS}from .${module}Template import *;")
+      set(ENV{ITK_STUB_PROXY_IMPORTS} "$ENV{ITK_STUB_PROXY_IMPORTS}from .${module}Proxy import *;")
     endif()
 
 

diff --git a/Wrapping/Generators/SwigInterface/igenerator.py b/Wrapping/Generators/SwigInterface/igenerator.py
@@ -200,7 +200,7 @@ def generate_class_pyi_def(
         )
 
     outputPYIMethodFile.write(
-        f"class _{class_name}Proxy({itk_class.parent_class}):\n"  # if
+        f"class {class_name}Proxy({itk_class.parent_class}):\n"  # if
     )
 
     if itk_class.is_enum and len(itk_class.enums) > 0:
@@ -647,6 +647,7 @@ def __init__(self, submoduleName, options, classes):
 
         # A dict of sets containing the .pyi python equivalent for all class methods and params
         self.classes = classes
+        self.python_parent_imports = []
         self.current_class = ""
 
         # a dict to let us use the alias name instead of the full c++ name. Without
@@ -1665,6 +1666,9 @@ def create_interfacefile(self, interfaceFile, idxFile, wrappersNamespace):
                     if ":" in base:
                         continue
 
+                    if base not in self.python_parent_imports:
+                        self.python_parent_imports.append(base)
+
                     if self.classes[self.current_class].parent_class == "":
                         self.classes[self.current_class].parent_class = f"_{base}Proxy"
 
@@ -1729,25 +1733,46 @@ def create_interfacefile(self, interfaceFile, idxFile, wrappersNamespace):
                 f.write(content)
 
 
-def init_submodule_pyi_file(pyiFile: Path, submodule_name: str) -> None:
+def init_submodule_pyi_template_file(pyiFile: Path, submodule_name: str) -> None:
+    with open(pyiFile, "w") as pyiFile:
+        pyiFile.write(
+            f"""# Interface for submodule: {submodule_name}
+from typing import Union, Any
+# additional imports
+from .support.template_class import itkTemplate as _itkTemplate
+
+\n"""
+        )
+
+
+def init_submodule_pyi_proxy_file(pyiFile: Path, submodule_name: str, parent_imports) -> None:
     with open(pyiFile, "w") as pyiFile:
         pyiFile.write(
-            f"""# Interface and Interface methods for submodule: {submodule_name}
+            f"""# Interface methods for submodule: {submodule_name}
 from typing import Union, Any
 # additional imports
 from .support.template_class import itkTemplate as _itkTemplate
+{ f"from ._proxies import {', '.join(parent_imports)}" if len(parent_imports) > 0 else ""}
 
 \n"""
         )
 
 
-def write_class_pyi(
-    pyiFile: Path, class_name: str, header_code: str, interfaces_code: str
+def write_class_template_pyi(
+        pyiFile: Path, class_name: str, header_code: str
 ) -> None:
-    # Write interface files to the stub directory to support editor autocompletions
+    # Write interface files to the stub directory to support editor autocompletion
     with open(pyiFile, "a+") as pyiFile:
         pyiFile.write(f"# Interface for class: {class_name}\n")
+        pyiFile.write(f"from ._proxies import {class_name}Proxy as _{class_name}Proxy\n")
         pyiFile.write(header_code)
+
+
+def write_class_proxy_pyi(
+        pyiFile: Path, class_name: str, interfaces_code: str
+) -> None:
+    # Write interface files to the stub directory to support editor autocompletion
+    with open(pyiFile, "a+") as pyiFile:
         pyiFile.write(f"# Interface methods for class: {class_name}\n")
         pyiFile.write(interfaces_code)
 
@@ -1958,6 +1983,8 @@ def generate_swig_input(submoduleName, classes):
             swig_input_generator.snakeCaseProcessObjectFunctions
         )
 
+        return swig_input_generator.python_parent_imports
+
     classes = {}
 
     ordered_submodule_list: List[str] = []
@@ -1971,11 +1998,15 @@ def generate_swig_input(submoduleName, classes):
     del submoduleNames
 
     for submoduleName in ordered_submodule_list:
-        generate_swig_input(submoduleName, classes)
+        parents = generate_swig_input(submoduleName, classes)
+        parent_imports = [f"{p}Proxy as _{p}Proxy" for p in parents]
         if options.pyi_dir != "":
-          init_submodule_pyi_file(
-            Path(f"{options.pyi_dir}/{submoduleName}.pyi"), submoduleName
-          )
+            init_submodule_pyi_template_file(
+                Path(f"{options.pyi_dir}/{submoduleName}Template.pyi"), submoduleName
+            )
+            init_submodule_pyi_proxy_file(
+                Path(f"{options.pyi_dir}/{submoduleName}Proxy.pyi"), submoduleName, parent_imports
+            )
 
     if options.pyi_dir != "":
       for itk_class in classes.keys():
@@ -1985,14 +2016,17 @@ def generate_swig_input(submoduleName, classes):
             outputPYIHeaderFile, outputPYIMethodFile, classes[itk_class]
         )
 
-        write_class_pyi(
-            Path(f"{options.pyi_dir}/{classes[itk_class].submodule_name}.pyi"),
+        write_class_template_pyi(
+            Path(f"{options.pyi_dir}/{classes[itk_class].submodule_name}Template.pyi"),
             itk_class,
             outputPYIHeaderFile.getvalue(),
+        )
+        write_class_proxy_pyi(
+            Path(f"{options.pyi_dir}/{classes[itk_class].submodule_name}Proxy.pyi"),
+            itk_class,
             outputPYIMethodFile.getvalue(),
         )
 
-
     snake_case_file = options.snake_case_file
     if len(snake_case_file) > 1:
         with open(snake_case_file, "w") as ff: