From 45894b66b0ae9e74037e6b4ddd2585be99e13497 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 30 Jun 2026 11:19:17 +0200 Subject: [PATCH] feat(cli): generate stubs in build dir when compiling --- midas/cli/commands/compile.py | 13 ++++++++---- midas/cli/commands/stubs.py | 10 +++++---- midas/generator/generator.py | 39 +++++++++++++++++++++++++++++------ tests/generator.py | 3 ++- 4 files changed, 50 insertions(+), 15 deletions(-) diff --git a/midas/cli/commands/compile.py b/midas/cli/commands/compile.py index 497a0f9..a59eea1 100644 --- a/midas/cli/commands/compile.py +++ b/midas/cli/commands/compile.py @@ -5,7 +5,7 @@ import sys from pathlib import Path -from typing import TextIO +from typing import Optional, TextIO import click @@ -19,18 +19,23 @@ from midas.utils import TypedAST @click.command(help="Compile source") @click.argument("file", type=click.File("r")) @click.option("-t", "--types", type=click.File("r"), multiple=True) +@click.option("-s", "--stubs", type=str, multiple=True) @click.option("--ignore-errors", is_flag=True) def compile( file: TextIO, types: tuple[TextIO], + stubs: tuple[str], ignore_errors: bool, ): source: str = file.read() source_path: Path = Path(file.name).resolve() checker = TypeChecker() - for types_file in types: - checker.import_midas(Path(types_file.name).resolve()) + type_files: list[tuple[Path, Optional[str]]] = [] + for i, types_file in enumerate(types): + in_path: Path = Path(types_file.name).resolve() + checker.import_midas(in_path) + type_files.append((in_path, stubs[i] if i < len(stubs) else None)) typed_ast: TypedAST = checker.type_check_source(source, str(source_path)) diagnostics: list[Diagnostic] = checker.diagnostics.copy() @@ -43,4 +48,4 @@ def compile( sys.exit(1) generator = Generator(workdir=source_path.parent, types=checker.types) - generator.generate(typed_ast, source_path) + generator.generate(typed_ast, source_path, type_files=type_files) diff --git a/midas/cli/commands/stubs.py b/midas/cli/commands/stubs.py index a5267c9..9eb028c 100644 --- a/midas/cli/commands/stubs.py +++ b/midas/cli/commands/stubs.py @@ -1,7 +1,7 @@ import ast import time from pathlib import Path -from typing import TextIO +from typing import Optional, TextIO import black import click @@ -38,15 +38,17 @@ class Handler(FileSystemEventHandler): @click.command(help="Generate stubs from Midas definitions") @click.argument("file", type=click.File("r")) -@click.option("-o", "--output", type=click.File("w"), default="-") +@click.option("-o", "--output", type=click.File("w")) @click.option("-w", "--watch", is_flag=True) def stubs( file: TextIO, - output: TextIO, + output: Optional[TextIO], watch: bool, ): source_path: Path = Path(file.name).resolve() - out_path: Path = Path(output.name).resolve() + out_path: Path = source_path.with_suffix(".pyi") + if output is not None: + out_path = Path(output.name).resolve() generate_stubs(source_path, out_path) if watch: diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 6386002..b136055 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -9,6 +9,7 @@ import midas.ast.midas as m import midas.ast.python as p from midas.ast.location import Location from midas.ast.printer import MidasPrinter +from midas.checker.checker import TypeChecker from midas.checker.registry import TypesRegistry from midas.checker.types import ( AliasType, @@ -30,6 +31,7 @@ from midas.checker.types import ( UnknownType, ) from midas.generator.constraints import ConstraintGenerator +from midas.generator.stubs import StubsGenerator from midas.utils import TypedAST @@ -64,8 +66,10 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): self.define_is_dataframe: bool = False self.define_is_column: bool = False - def generate_ast(self, typed_ast: TypedAST, src_path: Path) -> ast.AST: - self.rel_src_path = src_path.resolve().relative_to(self.workdir) + def set_src_path(self, path: Path): + self.rel_src_path = path.resolve().relative_to(self.workdir) + + def generate_ast(self, typed_ast: TypedAST) -> ast.AST: self._typed_ast = typed_ast body: list[ast.stmt] = self._visit_body(typed_ast.stmts) predicates: list[ast.stmt] = self._constraint_generator.get_definitions() @@ -83,10 +87,13 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): return module def generate( - self, typed_ast: TypedAST, src_path: Path, out_path: Optional[Path] = None + self, + typed_ast: TypedAST, + src_path: Path, + out_path: Optional[Path] = None, + type_files: Optional[list[tuple[Path, Optional[str]]]] = None, ) -> Path: - module: ast.AST = self.generate_ast(typed_ast, src_path) - compiled: str = ast.unparse(module) + self.set_src_path(src_path) if out_path is None: if self.build_dir.exists(): shutil.rmtree(self.build_dir) @@ -98,10 +105,30 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]): raise ValueError( f"Directory traversal, {self.rel_src_path} points outside of parent directory" ) - out_path.parent.mkdir(parents=True, exist_ok=True) + out_dir: Path = out_path.parent + out_dir.parent.mkdir(parents=True, exist_ok=True) + + if type_files is not None: + for in_path, out_name in type_files: + if out_name is None: + out_name = in_path.stem + self.generate_stubs(in_path, out_dir / f"{out_name}.py") + + module: ast.AST = self.generate_ast(typed_ast) + compiled: str = ast.unparse(module) + out_path.write_text(compiled) return out_path + def generate_stubs(self, in_path: Path, out_path: Path): + checker = TypeChecker() + checker.import_midas(in_path) + generator = StubsGenerator(checker.types) + module: ast.Module = generator.generate_stubs() + module = ast.fix_missing_locations(module) + output: str = ast.unparse(module) + out_path.write_text(output) + def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr: return ast.BinOp( left=expr.left.accept(self), diff --git a/tests/generator.py b/tests/generator.py index 067790f..f04992b 100644 --- a/tests/generator.py +++ b/tests/generator.py @@ -46,7 +46,8 @@ class GeneratorTester(Tester): if not any(d.type == DiagnosticType.ERROR for d in checker.diagnostics): generator = Generator(workdir=path.parent, types=checker.types) - result.compiled_ast = generator.generate_ast(typed_ast, path) + generator.set_src_path(path) + result.compiled_ast = generator.generate_ast(typed_ast) return result