fix(treeshake): make re-exports work (#39)

Author: jvllmr

rust/src/common/ast/providers/fully_qualified_name_provider.rs

@@ -31,7 +31,7 @@ impl FullyQualifiedNameProvider {
         }
     }
 
-    fn maybe_format_with_name_context(&self, name: &str) -> String {
+    pub fn resolve_qualified_name(&self, name: &str) -> String {
         if self.name_context.len() > 0 {
             format!("{}.{}", self.name_context, name)
         } else {
@@ -41,7 +41,7 @@ impl FullyQualifiedNameProvider {
 
     fn get_expr_qualified_name(&self, expr: &Expr) -> Option<String> {
         get_full_name_for_expr(expr).map(|name| match expr {
-            Expr::NamedExpr(_) => self.maybe_format_with_name_context(&name),
+            Expr::NamedExpr(_) => self.resolve_qualified_name(&name),
             _ => name,
         })
     }
@@ -58,14 +58,14 @@ impl FullyQualifiedNameProvider {
                 | Stmt::ClassDef(_)
                 | Stmt::FunctionDef(_)
                 | Stmt::AsyncFunctionDef(_) => {
-                    self.maybe_format_with_name_context(name)
+                    self.resolve_qualified_name(name)
                 }
                 _ => name.to_string(),
             })
             .collect()
     }
 
-    fn resolve_fully_qualified_name(&self, qualified_name: &str) -> Vec<String> {
+    pub fn resolve_fully_qualified_name(&self, qualified_name: &str) -> Vec<String> {
         let mut result: Vec<String> = Vec::new();
 
         for (key, value) in &self.imports_provider.active_imports {

rust/src/treeshake/references_counter.rs

@@ -153,21 +153,32 @@ impl ReferencesCounter {
         self.names_provider.name_context.len() == 0
     }
 
-    fn maybe_increase_stmt(&mut self, stmt: &Stmt) {
+    fn maybe_increase_stmt_selective<F>(&mut self, stmt: &Stmt, predicate: F)
+    where
+        F: Fn(&str) -> bool,
+    {
         for fqn in self.names_provider.get_stmt_fully_qualified_name(stmt) {
-            self.increase(&fqn);
+            if predicate(&fqn) {
+                self.increase(&fqn);
+            }
         }
 
         // bump for this node because it is a global name that could be imported somewhere else via star import
         if self.import_star_module_specs.len() > 0 {
             for module_spec in self.import_star_module_specs.to_owned() {
                 for full_name in get_full_name_for_stmt(stmt, &self.get_parent_package()) {
-                    self.increase(&format!("{}.{}", module_spec, full_name));
+                    if predicate(&full_name) {
+                        self.increase(&format!("{}.{}", module_spec, full_name));
+                    }
                 }
             }
         }
     }
 
+    fn maybe_increase_stmt(&mut self, stmt: &Stmt) {
+        self.maybe_increase_stmt_selective(stmt, |_| true);
+    }
+
     fn maybe_increase_expr(&mut self, expr: &Expr) {
         for fqn in self.names_provider.get_expr_fully_qualified_name(expr) {
             self.increase(&fqn);
@@ -283,13 +294,51 @@ impl Visitor for ReferencesCounter {
                     self.always_bump_context = true;
                 }
             }
-            Stmt::ImportFrom(import_from) => {
-                if import_from.names.len() == 1
-                    && import_from.names[0].name.as_str() == "*"
+            Stmt::Import(stmt_import) => {
+                // check if one of the names defined by this import was imported somewhere else
+                // if yes, bump reference of this import
+                for alias in &stmt_import.names {
+                    let defined_name = if let Some(alias_value) = &alias.asname {
+                        alias_value
+                    } else {
+                        &alias.name
+                    };
+                    for fqn in self.names_provider.resolve_fully_qualified_name(
+                        &self.names_provider.resolve_qualified_name(&defined_name),
+                    ) {
+                        if self.has_references_for_str(&fqn) {
+                            self.maybe_increase_stmt_selective(&stmt, |n| n == alias.name.as_str());
+                        }
+                    }
+                }
+            }
+
+            Stmt::ImportFrom(stmt_import_from) => {
+                // check if one of the names defined by this import was imported somewhere else
+                // if yes, bump reference of this import
+                for alias in &stmt_import_from.names {
+                    let defined_name = if let Some(alias_value) = &alias.asname {
+                        alias_value
+                    } else {
+                        &alias.name
+                    };
+                    for fqn in self.names_provider.resolve_fully_qualified_name(
+                        &self.names_provider.resolve_qualified_name(&defined_name),
+                    ) {
+                        if self.has_references_for_str(&fqn) {
+                            self.maybe_increase_stmt_selective(&stmt, |n| {
+                                n.ends_with(alias.name.as_str())
+                            });
+                        }
+                    }
+                }
+
+                if stmt_import_from.names.len() == 1
+                    && stmt_import_from.names[0].name.as_str() == "*"
                     && self.module_spec_has_references()
                 {
                     if let Ok(module_specs) = get_import_from_absolute_module_spec(
-                        &import_from,
+                        &stmt_import_from,
                         &self.get_parent_package(),
                     ) {
                         self.import_star_module_specs.extend(module_specs);

tests/test_treeshake/packages/re_exports/__init__.py

@@ -0,0 +1,6 @@
+from .hello_world import hello_world
+from .hello_world import inner_hello_world_alias
+
+def main() -> None:
+    hello_world()
+    inner_hello_world_alias.hello_world()

tests/test_treeshake/packages/re_exports/__main__.py

@@ -0,0 +1,5 @@
+from . import main
+
+
+if __name__ == "__main__":
+    main()

tests/test_treeshake/packages/re_exports/hello_world/__init__.py

@@ -0,0 +1,5 @@
+from .inner_hello_world import hello_world, useless_func
+from .moin_world import moin_world
+import re_exports.hello_world.inner_hello_world as inner_hello_world_alias, re_exports.hello_world.moin_world as moin_world_alias  # type: ignore[import-not-found]
+
+__all__ = ["hello_world", "inner_hello_world_alias", "moin_world", "moin_world_alias", "useless_func"]

tests/test_treeshake/packages/re_exports/hello_world/inner_hello_world/__init__.py

@@ -0,0 +1,5 @@
+def hello_world() -> None:
+    print("Hello World!")
+
+def useless_func() -> None:
+    pass

tests/test_treeshake/packages/re_exports/hello_world/moin_world/__init__.py

@@ -0,0 +1,2 @@
+def moin_world() -> None:
+    print("Moin Welt!")

tests/test_treeshake/test_treeshake_package.py

@@ -78,3 +78,39 @@ def test_treeshake_package_preserve_with_decorators(
     init_file_content = init_file.read_text()
     assert "@dataclass\nclass MyClass:\n    pass" in init_file_content
     assert "@contextmanager\ndef my_context_manager() ->" in init_file_content
+
+
+def test_treeshake_package_re_exports(
+    run_treeshake_package: RunTreeshakePackageT,
+) -> None:
+    source_path = TEST_PACKAGES_DIR / "re_exports"
+    result_path = run_treeshake_package(source_path)
+
+    hello_world_init = result_path / "hello_world" / "__init__.py"
+    assert hello_world_init.exists()
+    hello_world_init_content = hello_world_init.read_text()
+    assert "from .inner_hello_world import hello_world" in hello_world_init_content
+    assert (
+        "import re_exports.hello_world.inner_hello_world as inner_hello_world_alias"
+        in hello_world_init_content
+    )
+    assert (
+        "re_exports.hello_world.moin_world as moin_world_alias"
+        not in hello_world_init_content
+    )
+    assert "useless_func" not in hello_world_init_content
+    assert "from .moin_world import moin_world" not in hello_world_init_content
+
+    assert not (result_path / "hello_world" / "moin_world").exists()
+
+    inner_hello_world_init = hello_world_init = (
+        result_path / "hello_world" / "inner_hello_world" / "__init__.py"
+    )
+    assert inner_hello_world_init.exists()
+    inner_hello_world_init_content = inner_hello_world_init.read_text()
+    assert (
+        "def hello_world() -> None:\n    print('Hello World!')"
+        in inner_hello_world_init_content
+    )
+
+    assert "def useless_func() -> None:" not in inner_hello_world_init_content