Enhance readability following merge PR #27 (Callback to customize the…

… base classes used in generated bindings)

src/litgen/__init__.py

@@ -16,6 +16,7 @@
 from litgen.litgen_generator import (
     LitgenGenerator,
     GeneratedCodes,
+    GeneratedCodeType,
     write_generated_code_for_files,
     write_generated_code_for_file,
     generate_code,
@@ -37,6 +38,7 @@
     # When it is needed to have different options per c++ header file
     "LitgenGenerator",
     "GeneratedCodes",
+    "GeneratedCodeType",
     "generate_code_for_file",
     # Configure replacements
     "standard_type_replacements",

src/litgen/internal/adapted_types/adapted_class.py

@@ -470,12 +470,14 @@ def stub_lines(self) -> list[str]:
         def str_parent_classes_python() -> str:
             parents: list[str] = []
             custom_derived = (
-                [] if not self.options.class_base_custom_derivation__callback 
-                    else self.options.class_base_custom_derivation__callback(self, True))
+                []
+                if not self.options.class_custom_inheritance__callback
+                else self.options.class_custom_inheritance__callback(self, litgen.GeneratedCodeType.stub)
+            )
 
             if not custom_derived and not self.cpp_element().has_base_classes():
                 return ""
-            
+
             if custom_derived:
                 for custom_base in custom_derived:
                     parents.append(custom_base)
@@ -614,8 +616,10 @@ def make_pyclass_creation_code() -> str:
             # fill py::class_ additional template params (base classes, nodelete, etc)
             other_template_params_list = []
             custom_derived = (
-                [] if not self.options.class_base_custom_derivation__callback 
-                    else self.options.class_base_custom_derivation__callback(self, False))
+                []
+                if not self.options.class_custom_inheritance__callback
+                else self.options.class_custom_inheritance__callback(self, litgen.GeneratedCodeType.pydef)
+            )
 
             if custom_derived:
                 for custom_base in custom_derived:

src/litgen/litgen_generator.py

@@ -3,6 +3,7 @@
 import os
 import subprocess
 from dataclasses import dataclass
+from enum import Enum
 
 from codemanip import code_utils
 
@@ -61,6 +62,12 @@ def add_python_exe_folder_to_env_path() -> None:
     return _apply_black_formatter_pyi_via_subprocess(options, file)
 
 
+class GeneratedCodeType(Enum):
+    pydef = 1
+    stub = 2
+    glue = 3
+
+
 @dataclass
 class _GeneratedCode:
     source_filename: CppFilename

src/litgen/options.py

@@ -12,7 +12,8 @@
 from litgen.internal.class_iterable_info import ClassIterablesInfos
 
 if TYPE_CHECKING:
-    from litgen.internal.adapted_types import AdaptedFunction
+    from litgen.internal.adapted_types import AdaptedFunction, AdaptedClass
+    from litgen.litgen_generator import GeneratedCodeType
 
 
 class BindLibraryType(Enum):
@@ -505,6 +506,13 @@ class LitgenOptions:
     # - [Understanding Holder Types in pybind11](https://pybind11.readthedocs.io/en/stable/advanced/classes.html#custom-smart-pointers)
     class_held_as_shared__regex: str = ""
 
+    # class_custom_inheritance__callback:
+    # (advanced) A callback to customize the base classes used in generated bindings.
+    # The first parameter is the AdaptedClass, representing the C++ class being adapted.
+    # The second parameter is the GeneratedCodeType, indicating whether stub or pydef code is being generated.
+    # An example usage can be found in: src/litgen/tests/option_class_custom_inheritance__callback_test.py
+    class_custom_inheritance__callback: Callable[[AdaptedClass, GeneratedCodeType], list[str]] | None = None
+
     # ------------------------------------------------------------------------------
     # Templated class options
     # ------------------------------------------------------------------------------
@@ -531,11 +539,6 @@ class LitgenOptions:
     # the generated classes together
     class_template_decorate_in_stub: bool = True
 
-    # This callback Callback to customize the base classes used in generated bindings
-    # First param is the AdoptedClass
-    # Second indicates context - True for python stub, False for CPP bindings
-    class_base_custom_derivation__callback: Callable[[Any, bool], list[str]] | None = None
-
     # ------------------------------------------------------------------------------
     # Adapt class members
     # ------------------------------------------------------------------------------

src/litgen/tests/option_custom_class_derive_test.py → src/litgen/tests/option_class_custom_inheritance__callback_test.py

@@ -1,26 +1,18 @@
 from __future__ import annotations
 from codemanip import code_utils
 import litgen
+from litgen.internal.adapted_types import AdaptedClass
+from litgen import GeneratedCodeType
 
-def test_custom_classes_base_option():
-    """
-    Example of how the callback mechanism could be used in practice to handle reference return policies
-    """
-
-    def handle_classes_base(cls, for_python_stub):
-
-        bases = []
-
-        elem = cls.cpp_element()
-
-        if elem.class_name == "SecondClass":
-            bases.append("FirstClass" if for_python_stub else "CustomNS::FirstClass")
 
-        return bases
-
+def test_class_custom_inheritance__callback():
+    """Example of how the callback mechanism `options.class_custom_inheritance__callback`
+    can be used in practice to add a base class
     """
-    ## First class is in another file, and won't be handled by litgen for another files
 
+    # Let's suppose that we have the following C++ code in another file,
+    # which was not processed by litgen (or not yet)
+    """
         namespace CustomNS {
             class FirstClass {
             public:
@@ -31,6 +23,7 @@ class FirstClass {
         }
     """
 
+    # And we are processing the following code:
     code = """
         class SecondClass : CustomNS::FirstClass {
         public:
@@ -47,8 +40,17 @@ class ThirdClass : SecondClass {
         };
     """
 
+    # This will be our callback to add the base class: it returns the base class which we should add
+    # (with a syntax that depends slightly on the generated code type)
+    def handle_classes_base(cls: AdaptedClass, generated_code_type: GeneratedCodeType) -> list[str]:
+        bases = []
+        elem = cls.cpp_element()
+        if elem.class_name == "SecondClass":
+            bases.append("FirstClass" if generated_code_type == GeneratedCodeType.stub else "CustomNS::FirstClass")
+        return bases
+
     options = litgen.LitgenOptions()
-    options.class_base_custom_derivation__callback = handle_classes_base
+    options.class_custom_inheritance__callback = handle_classes_base
     generated_code = litgen.generate_code(options, code)
 
     code_utils.assert_are_codes_equal(