Gen stubs

Author: jacobtomlinson

ci/gen_stubs.py

@@ -2,54 +2,103 @@
 # SPDX-License-Identifier: BSD 3-Clause License
 # Update sync code stubs
 import ast
+import subprocess
 
-import kr8s._objects as base_objects
+ELLIPSIS_NODE = ast.parse("...").body
+# TODO: Preserve docstrings in method/function definitions
+# TODO: Remove unused imports
 
+# Load the AST for the base object definitions
 with open("kr8s/_objects.py") as base_source_fh:
     base_source = base_source_fh.read()
     base_tree = ast.parse(base_source)
 
-
+# Load the AST for the sync API
 with open("kr8s/objects.py") as source_fh:
     source = source_fh.read()
     tree = ast.parse(source)
 
-    # Update module docstring
-    if isinstance(tree.body[0], ast.Expr) and isinstance(
-        tree.body[0].value, ast.Constant
-    ):
-        tree.body[0].value.value += (
-            "\n"
-            "WARNING: This typing stub file was generated by ci/gen_stubs.py\n"
-            "DO NOT MODIFY DIRECTLY\n"
-        )
-
-    for node in tree.body:
-        if isinstance(node, ast.ImportFrom):
-            if node.level != 1:
-                tree.body.remove(node)
-
-        elif isinstance(node, ast.ClassDef):
-            # Clear the sync decorator
-            for decorator in node.decorator_list:
-                if isinstance(decorator, ast.Name) and decorator.id == "sync":
-                    node.decorator_list.remove(decorator)
-            node.body = [subnode for subnode in node.body if not isinstance(subnode, ast.Assign)]
-            # Populate public methods from base classes
-            for base in node.bases:
-                if isinstance(base, ast.Name):
-                    # import the base class in this script with importlib
-                    base_class = getattr(base_objects, base.id[1:])
-                    # iterate over the methods of the base class
-                    for base_class_attr in base_class.__dict__.keys():
-                        method = getattr(base_class, base_class_attr)
-                        if callable(method):
-                            if not base_class_attr.startswith("_") and not base_class_attr.startswith("async_"):
-                                node.body.append(ast.parse(f"def {base_class_attr}(self): ...").body[0])
-                                
-
-        else:
-            pass
-
-    with open("kr8s/objects_tmp.pyi", "w") as target_fh:
-        target_fh.write(ast.unparse(tree))
+# Update module docstring
+if isinstance(tree.body[0], ast.Expr) and isinstance(tree.body[0].value, ast.Constant):
+    tree.body[0].value.value += (
+        "\n"
+        "WARNING: This typing stub file was generated by ci/gen_stubs.py\n"
+        "DO NOT MODIFY DIRECTLY\n"
+    )
+
+
+# Add all imports from base_tree (these will be pruned later)
+for node in base_tree.body:
+    if isinstance(node, ast.Import):
+        tree.body.insert(2, node)
+    elif isinstance(node, ast.ImportFrom):
+        # If not __future__ import, add it
+        if node.module != "__future__":
+            tree.body.insert(2, node)
+
+for node_index, node in enumerate(tree.body):
+
+    # # Remove non-relative imports
+    # if isinstance(node, ast.ImportFrom):
+    #     if node.level != 1:
+    #         tree.body.remove(node)
+
+    # Update class definitions
+    if isinstance(node, ast.ClassDef):
+        # Replace the class definitions with the bass class
+        for base_node in base_tree.body:
+            if isinstance(base_node, ast.ClassDef) and base_node.name == node.name:
+                if node.name == "APIObject":
+                    base_node.bases = node.bases
+
+                # Replace all methods with type stubs
+                for subnode_index, subnode in enumerate(base_node.body):
+                    if isinstance(subnode, ast.FunctionDef):
+                        subnode.body = ELLIPSIS_NODE
+                    elif isinstance(subnode, ast.AsyncFunctionDef):
+                        base_node.body[subnode_index] = ast.FunctionDef(
+                            name=subnode.name,
+                            args=subnode.args,
+                            body=ELLIPSIS_NODE,
+                            decorator_list=subnode.decorator_list,
+                            returns=subnode.returns,
+                            type_comment=subnode.type_comment,
+                            type_params=subnode.type_params,
+                            lineno=subnode.lineno,
+                        )
+                        pass
+                    elif isinstance(subnode, ast.AnnAssign):
+                        # Remove value
+                        subnode.value = None
+                    else:
+                        pass
+
+                tree.body[node_index] = base_node
+
+    elif isinstance(node, ast.FunctionDef):
+        # Replace all functions with version from bass_tree
+        for base_node in base_tree.body:
+            if (
+                isinstance(base_node, ast.FunctionDef)
+                or isinstance(base_node, ast.AsyncFunctionDef)
+            ) and base_node.name == node.name:
+                # Replace the function definition
+                tree.body[node_index] = ast.FunctionDef(
+                    name=base_node.name,
+                    args=base_node.args,
+                    body=ELLIPSIS_NODE,
+                    decorator_list=base_node.decorator_list,
+                    returns=base_node.returns,
+                    type_comment=base_node.type_comment,
+                    type_params=base_node.type_params,
+                    lineno=base_node.lineno,
+                )
+    else:
+        pass
+
+with open("kr8s/objects.pyi", "w") as target_fh:
+    target_fh.write(ast.unparse(tree))
+
+# Run black and ruff
+subprocess.run(["black", "kr8s/objects.pyi"])
+subprocess.run(["ruff", "check", "--fix", "kr8s/objects.pyi"])

kr8s/objects.py

@@ -5,6 +5,8 @@
 This module provides classes that represent Kubernetes resources.
 These classes are used to interact with resources in the Kubernetes API server.
 """
+from __future__ import annotations
+
 from functools import partial
 
 from ._async_utils import run_sync, sync
@@ -347,8 +349,21 @@ class Table(_Table):
     _asyncio = False
 
 
-object_from_name_type = run_sync(partial(_object_from_name_type, _asyncio=False))
-objects_from_files = run_sync(partial(_objects_from_files, _asyncio=False))
-get_class = partial(_get_class, _asyncio=False)
-new_class = partial(_new_class, asyncio=False)
-object_from_spec = partial(_object_from_spec, _asyncio=False)
+def object_from_name_type(*args, **kwargs):
+    return run_sync(partial(_object_from_name_type, _asyncio=False))(*args, **kwargs)
+
+
+def objects_from_files(*args, **kwargs):
+    return run_sync(partial(_objects_from_files, _asyncio=False))(*args, **kwargs)
+
+
+def get_class(*args, **kwargs):
+    return partial(_get_class, _asyncio=False)(*args, **kwargs)
+
+
+def new_class(*args, **kwargs):
+    return partial(_new_class, asyncio=False)(*args, **kwargs)
+
+
+def object_from_spec(*args, **kwargs):
+    return partial(_object_from_spec, _asyncio=False)(*args, **kwargs)