fix(bundle): transform annotation literals (#40)

Author: jvllmr

Cargo.lock

@@ -12,4 +12,4 @@ crate-type = ["cdylib"]
 pyo3 = { version = "0.23.4", features = ["extension-module"] }
 rustpython-ast = { version = "0.4.0", features = ["visitor"] }
 rustpython-parser = { version = "0.4.0" }
-rustpython-unparser = { version = "0.2.0", features = ["transformer"] }
+rustpython-unparser = { version = "0.2.1", features = ["transformer"] }

rust/Cargo.toml

@@ -1,8 +1,8 @@
 use crate::common::module_spec::{get_top_level_package, is_in_std_lib};
 use pyo3::pyfunction;
 use rustpython_ast::{
-    text_size::TextRange, Alias, Expr, ExprAttribute, ExprName, Identifier, StmtImport,
-    StmtImportFrom, Suite,
+    text_size::TextRange, Alias, Constant, Expr, ExprAttribute, ExprConstant, ExprName, Identifier,
+    StmtImport, StmtImportFrom, Suite,
 };
 use rustpython_parser::Parse;
 use rustpython_unparser::{transformer::Transformer, Unparser};
@@ -12,6 +12,7 @@ struct ImportsTransformer {
     top_level_package: String,
     vendor_module_name: String,
     affected_names: HashSet<String>,
+    is_in_annotation: bool,
 }
 
 impl ImportsTransformer {
@@ -20,6 +21,7 @@ impl ImportsTransformer {
             top_level_package,
             vendor_module_name,
             affected_names: HashSet::new(),
+            is_in_annotation: false,
         }
     }
 
@@ -51,7 +53,37 @@ impl ImportsTransformer {
 }
 
 impl Transformer for ImportsTransformer {
-    fn generic_visit_stmt_import(&mut self, stmt: StmtImport) -> Option<StmtImport> {
+    fn on_enter_annotation(&mut self, _: &Expr) {
+        self.is_in_annotation = true;
+    }
+
+    fn on_exit_annotation(&mut self, _: &Option<Expr>) {
+        self.is_in_annotation = false;
+    }
+
+    fn visit_expr_constant(&mut self, mut expr: ExprConstant) -> Option<ExprConstant> {
+        if self.is_in_annotation {
+            match &expr.value {
+                Constant::Str(str_value) => {
+                    if str_value.contains(".") {
+                        let name_parts: Vec<&str> = str_value.splitn(2, ".").collect();
+                        let module_part = name_parts[0];
+                        if self.affected_names.contains(module_part) {
+                            expr.value =
+                                Constant::Str(format!("{}.{}", self.get_vendor_string(), str_value))
+                        }
+                    }
+
+                    Some(expr)
+                }
+                _ => self.generic_visit_expr_constant(expr),
+            }
+        } else {
+            self.generic_visit_expr_constant(expr)
+        }
+    }
+
+    fn visit_stmt_import(&mut self, stmt: StmtImport) -> Option<StmtImport> {
         Some(StmtImport {
             names: stmt
                 .names

rust/src/bundle/imports_transformer.rs

@@ -0,0 +1,10 @@
+import typer
+import typer as typerino
+
+random_literal = "typer.Typer"
+
+def modify_app(app: "typer.Typer") -> "typer.Typer":
+    return app
+
+def modify_app2(app2: "typerino.Typer") -> "typerino.Typer":
+    return app2

tests/test_bundle/packages/annotation_string_literals/__init__.py

@@ -121,7 +121,9 @@ def test_bundle_relative_imports(run_bundle_package: RunBundlePackageT) -> None:
     assert (result_path / "module" / "hello_world.py").exists()
 
 
-def test_correct_transformer_recursion(run_bundle_package: RunBundlePackageT) -> None:
+def test_bundle_correct_transformer_recursion(
+    run_bundle_package: RunBundlePackageT,
+) -> None:
     _, result_path = run_bundle_package(
         "correct_transformer_recursion", "correct_transformer_recursion"
     )
@@ -133,3 +135,26 @@ def test_correct_transformer_recursion(run_bundle_package: RunBundlePackageT) ->
         "correct_transformer_recursion._vendor.typer.style('Hello World!', fg=correct_transformer_recursion._vendor.typer.colors.BRIGHT_MAGENTA)"
         in init_file_content
     )
+
+
+def test_bundle_annotation_string_literals(
+    run_bundle_package: RunBundlePackageT,
+) -> None:
+    _, result_path = run_bundle_package(
+        "annotation_string_literals", "annotation_string_literals"
+    )
+
+    init_file = result_path / "__init__.py"
+    init_file_content = init_file.read_text()
+
+    assert "random_literal = 'typer.Typer'" in init_file_content
+
+    assert (
+        "def modify_app(app: 'annotation_string_literals._vendor.typer.Typer') -> 'annotation_string_literals._vendor.typer.Typer':\n    return app"
+        in init_file_content
+    )
+
+    assert (
+        "def modify_app2(app2: 'typerino.Typer') -> 'typerino.Typer':\n    return app2"
+        in init_file_content
+    )