typegen 0.6.0 (#86)

.github/workflows/python-app.yml

@@ -34,7 +34,7 @@ jobs:
 
     - name: Install Ruff 
       run: |
-        pipx install ruff==0.1.5
+        pipx install ruff==0.3.1
 
     - name: Lint Test 
       run: |

README.md

@@ -1,9 +1,5 @@
 # Polugins
 
-> [!CAUTION]
-> The type generation part is not working for new versions until I implement https://github.com/StefanBRas/polugins/issues/80.
-> If you need this, feel free to comment on the issue or create a new one. This is kinda low on my priorities until it's actually needed by someone.
-
 Minimal "plugin" system for polars.
 
 With Polugins:
@@ -45,7 +41,7 @@ class MyNamespace:
 
 You also use an external package `example-package` that exposes a `LazyFrame` namespace called `external`.
 
-in your `pyproject.toml` add a section:
+In your `pyproject.toml` add a section:
 
 ```toml
 [tool.polugins.lazyframe]
@@ -78,9 +74,9 @@ You need to make sure that you have called `register_namespaces` before trying t
 As an alternative, polars is re-exported through `polugins` such that entrypoint, configuration and environment variable namespaces are automagically registered:
 
 ```python
-from polugins import pl
+from polugins import pl # Namespaces are now registered
+
 
-# Namespaces are now registered
 (
   pl.LazyFrame()
   .external.some_method()
@@ -91,7 +87,6 @@ from polugins import pl
 Since the registration is dynamic, your linter, typechecker and IDE will not know about these namespaces.
 To solve this, install `polugins_type_gen` and then run `polugins stubs` from the root of you package. Type information is now generated and should be automatically picked up by your tools. 
 
-
 ## Usage
 
 Namespaces can be registered in three ways:
@@ -139,7 +134,7 @@ Which will register a  `LazyFrame` namespace located at `my_package.namespaces`
 
 To generate types install the python package `polugins_type_gen` and then run `polugins stubs` to create type stubs at "./typings".
 
-`polugins_type_gen` is only used as a CLI tool so it's recommended to put this in developer dependencies or installed with a tool like `pipx`.
+`polugins_type_gen` is only used as a CLI tool so it's recommended to put this in developer dependencies or run with `uvx`/`pipx`.
 
 ## Third party Package example
 
@@ -169,6 +164,5 @@ Don't use the `pl.api.register_x` in your package. This will make the extension
 Just a thin wrapper around `polars.api.register_x_namespace` and then using `importlib.metadata` to collect
 namespaces from external packages.
 
-Types are generated by using mypy to create stubs for lazyframe, dataframe, expr and series and then adding the
-namespaces to these type stubs.
+Types are generated by reading in the source files for the polars classes, strip all functions bodies and then add the functions from the namespaces.
 

polugins/src/polugins/_types.py

@@ -4,9 +4,7 @@
 from typing import Callable
 
 import polars as pl
-from typing_extensions import Self, TypeVar
-
-NS = TypeVar("NS")
+from typing_extensions import Self
 
 
 class ExtensionClass(Enum):

polugins_type_gen/CHANGELOG.md

@@ -2,8 +2,12 @@
 
 ## Unreleased
 
+## 6.0.0
+
+- Generate stubs from ast and not with mypy.
+- Allow setting output path
+- Allow finding polars in current venv or given path.
 - Fix nested module paths in stubs.
-- Generate stubs from ast and not with mypy
 
 ## 0.5.8
 

polugins_type_gen/Taskfile.yml

@@ -5,10 +5,15 @@ tasks:
     cmds:
       - poetry run python scripts/generate_polars_stubs.py
 
+  install:
+    cmds:
+      - poetry install --with dev
+
   test:
-    deps: ['nox']
+    deps:
+      - install
     cmds:
-      - poetry run nox -r -s test 
+      - poetry run pytest
 
   format:
     cmds:
@@ -18,8 +23,3 @@ tasks:
   lint:
     cmds:
       - poetry run ruff check .
-
-  nox:
-    sources: ["pyproject.toml", "poetry.lock"]
-    cmds:
-      - poetry install --only=nox

polugins_type_gen/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "polugins_type_gen"
-version = "0.6.0rc1"
+version = "0.6.0"
 description = "Type stub generator Polugins."
 authors = ["StefanBRas <opensource@bruhn.io>"]
 readme = "README.md"
@@ -15,7 +15,6 @@ optional = true
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"
-nox = "^2023.4.22"
 coverage = "^7.2.7"
 pytest-env = "^0.8.2"
 pytest-cov = "^4.1.0"

polugins_type_gen/src/polugins_type_gen/_ast.py

@@ -65,10 +65,23 @@ def parse(module_source: str):
     return ast.unparse(parsed)
 
 
+def parse_from_venv(venv_path: Path, extension_class: ExtensionClass) -> ast.Module:
+    venv_libs = list((venv_path / "lib").glob("python*"))
+    assert len(venv_libs) > 0, "No python versions in venv found."
+    if len(venv_libs) > 1:
+        print(f"Warning: Multiple python versions in venv found. Using the first one ({venv_libs[0].name})")
+    module_import_path = venv_libs[0] / "site-packages" / extension_class.import_path
+    return parse_polars_module(module_import_path.with_suffix(".py").read_text())
+
+
 def parse_from_current_env(extension_class: ExtensionClass) -> ast.Module:
     module_import_path = ".".join(_IMPORT_PATHS[extension_class].parts)
     module = importlib.import_module(module_import_path)
     module_source = inspect.getsource(module)
+    return parse_polars_module(module_source)
+
+
+def parse_polars_module(module_source: str) -> ast.Module:
     parsed = ast.parse(module_source, type_comments=False)
     return RemoveBodies().visit(parsed)
 

polugins_type_gen/src/polugins_type_gen/cli.py

@@ -1,11 +1,17 @@
+import argparse
 import ast
 import importlib.resources as importlib_resources
-import sys
 from pathlib import Path
 
 from polugins.main import _get_namespaces  # TODO: makes this non private
+from typing_extensions import Optional
 
-from polugins_type_gen._ast import parse_from_current_env as _parse_from_current_env
+from polugins_type_gen._ast import (
+    parse_from_current_env as _parse_from_current_env,
+)
+from polugins_type_gen._ast import (
+    parse_from_venv as _parse_from_venv,
+)
 
 
 class NoNamespaceRegisteredException(Exception):
@@ -17,8 +23,8 @@ def has_version(version: str) -> bool:
     return (files / "_stubs" / version).is_dir()
 
 
-def create_stubs():
-    output_dir = Path("typings")
+def create_stubs(venv_path: Optional[Path] = None, output_dir: Optional[Path] = None):
+    output_dir = output_dir or Path("typings")
     all_namespaces = _get_namespaces()
 
     if all(namespace == {} for namespace in all_namespaces.values()):
@@ -27,7 +33,13 @@ def create_stubs():
 
     for extension_class, namespaces in all_namespaces.items():
         if namespaces:
-            stub_ast = _parse_from_current_env(extension_class)
+            print(f"Found namespaces for {extension_class.value}:")
+            for name, namespace in namespaces.items():
+                print(f"  {name}: {namespace}")
+            if venv_path:
+                stub_ast = _parse_from_venv(venv_path, extension_class)
+            else:
+                stub_ast = _parse_from_current_env(extension_class)
             new_class_nodes = []
             module_imports = set()
             for node in stub_ast.body:
@@ -54,18 +66,50 @@ def create_stubs():
             output_path = output_dir / extension_class.import_path
             output_path.parent.mkdir(exist_ok=True, parents=True)
             output_path.with_suffix(".pyi").write_text(ast.unparse(stub_ast))
+            print(f"Generated stubs for {extension_class.value} in {output_path.with_suffix('.pyi')}")
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser(
+        prog="Polugins Type Gen",
+        description="Generates type stubs for polars with added namespaces.",
+    )
+    parser.add_argument(
+        "command",
+        help="The command to run",
+        choices=["stubs", "version"],
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        help="The input venv to find Polars in. If not selected, uses currently activated venv.",
+        required=False,
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        help="The output directory for the generated stubs. Default is ./typings/",
+        default="typings",
+        required=False,
+    )
+    return parser
 
 
 def cli():
-    if len(sys.argv) == 1:
+    parser = get_argparser()
+    args = parser.parse_args()
+    if args.command == "stubs":
+        create_stubs(
+            venv_path=Path(args.input) if args.input else None,
+            output_dir=Path(args.output) if args.output else None,
+        )
+    elif args.command == "version":
         from polugins_type_gen._version import __version__
 
         print(f"Polugins Type Gen version: {__version__}")
-    elif sys.argv[1] == "stubs":
-        print("generating stubs at ./typings/")
-        create_stubs()
     else:
-        print("Use `polugins stubs` to generate type stubs.")
+        msg = "Unknown command. Use `polugins stubs` to generate type stubs."
+        raise ValueError(msg)
 
 
 if __name__ == "__main__":