fix(treeshake): default arg values are visited

RustPython/Parser#133

Author: jvllmr

rust/src/common/ast/mod.rs

@@ -1,6 +1,6 @@
 pub mod full_name;
 pub mod providers;
-
+pub mod visitor_patch;
 use pyo3::{
     exceptions::{PyImportError, PyValueError},
     PyResult,

rust/src/common/ast/visitor_patch.rs

@@ -0,0 +1,41 @@
+use rustpython_ast::{Arg, ArgWithDefault, Arguments, Visitor};
+// patch for missing visits in RustPython visitor
+// https://github.com/RustPython/Parser/issues/133
+pub trait VisitorPatch: Visitor {
+    fn generic_visit_arg_patch(&mut self, arg: Arg) {
+        if let Some(annotation) = arg.annotation {
+            self.visit_expr(*annotation);
+        }
+    }
+
+    fn generic_visit_arguments_patch(&mut self, args: Arguments) {
+        for arg in args.args {
+            self.visit_arg_with_default(arg);
+        }
+
+        for posonly in args.posonlyargs {
+            self.visit_arg_with_default(posonly);
+        }
+
+        for kwonly in args.kwonlyargs {
+            self.visit_arg_with_default(kwonly);
+        }
+        if let Some(vararg) = args.vararg {
+            self.visit_arg(*vararg);
+        }
+        if let Some(kwarg) = args.kwarg {
+            self.visit_arg(*kwarg);
+        }
+    }
+
+    fn visit_arg_with_default(&mut self, arg: ArgWithDefault) {
+        self.generic_visit_arg_with_default(arg);
+    }
+
+    fn generic_visit_arg_with_default(&mut self, arg: ArgWithDefault) {
+        self.visit_arg(arg.def);
+        if let Some(default) = arg.default {
+            self.visit_expr(*default);
+        }
+    }
+}

rust/src/treeshake/references_counter.rs

@@ -5,14 +5,17 @@ use std::{
 };
 
 use pyo3::{pyclass, pymethods};
-use rustpython_ast::{Expr, ExprCompare, Stmt, StmtImport, StmtImportFrom, Suite, Visitor};
+use rustpython_ast::{
+    Arg, Arguments, Expr, ExprCompare, Stmt, StmtImport, StmtImportFrom, Suite, Visitor,
+};
 use rustpython_parser::Parse;
 
 use crate::common::{
     ast::{
         full_name::{get_full_name_for_expr, get_full_name_for_stmt},
         get_import_from_absolute_module_spec,
         providers::fully_qualified_name_provider::FullyQualifiedNameProvider,
+        visitor_patch::VisitorPatch,
     },
     module_spec::get_parent_package,
 };
@@ -372,6 +375,13 @@ impl Visitor for ReferencesCounter {
 
         match expr {
             Expr::Call(_) => {
+                if let Some(full_name) = get_full_name_for_expr(&expr) {
+                    println!(
+                        "Default {} {}",
+                        full_name,
+                        self.is_global_scope() && self.module_spec_has_references()
+                    )
+                }
                 if self.is_global_scope() && self.module_spec_has_references() {
                     self.maybe_increase_expr(&expr);
                     self.always_bump_context = true;
@@ -395,4 +405,14 @@ impl Visitor for ReferencesCounter {
         self.names_provider.visit_import_from(&node);
         self.generic_visit_stmt_import_from(node);
     }
+
+    fn generic_visit_arg(&mut self, node: Arg) {
+        self.generic_visit_arg_patch(node);
+    }
+
+    fn generic_visit_arguments(&mut self, node: Arguments) {
+        self.generic_visit_arguments_patch(node);
+    }
 }
+
+impl VisitorPatch for ReferencesCounter {}

tests/test_treeshake/packages/param_default_value/__init__.py

@@ -0,0 +1,22 @@
+import typing as t
+import dataclasses
+
+@dataclasses.dataclass
+class DefaultPlaceholder:
+    value: t.Any
+
+def Default(value: t.Any) -> DefaultPlaceholder:
+    return DefaultPlaceholder(value=value)
+
+
+class DefaultUser:
+    def __init__(self, value: t.Any = Default(42)) -> None:
+        self.value = value
+
+
+def main() -> None:
+    user = DefaultUser()
+    print(user.value)
+
+if __name__ == "__main__":
+    main()

tests/test_treeshake/test_treeshake_package.py

@@ -114,3 +114,16 @@ def test_treeshake_package_re_exports(
     )
 
     assert "def useless_func() -> None:" not in inner_hello_world_init_content
+
+
+def test_treeshake_package_param_default_value(
+    run_treeshake_package: RunTreeshakePackageT,
+) -> None:
+    source_path = TEST_PACKAGES_DIR / "param_default_value"
+    result_path = run_treeshake_package(source_path)
+    init_file = result_path / "__init__.py"
+    init_file_content = init_file.read_text()
+    assert (
+        "def Default(value: t.Any) -> DefaultPlaceholder:\n    return DefaultPlaceholder(value=value)"
+        in init_file_content
+    )