Repository: SuadeLabs/rattr (language: python)

Issue:
Allow anonymouse functions in calls to `defaultdict`
The same as #5 but with `defaultdict`.


Discussion / hints:

Fix patch:
diff --git a/ratter/plugins/__init__.py b/ratter/plugins/__init__.py
index 4bf424f8..77b27ea9 100644
--- a/ratter/plugins/__init__.py
+++ b/ratter/plugins/__init__.py
@@ -7,7 +7,7 @@
     CustomFunctionHandler,
 )
 from ratter.analyser.util import MODULE_BLACKLIST_PATTERNS
-from ratter.plugins.analysers import BUILTIN_FUNCTION_ANALYSERS as BUILTINS
+from ratter.plugins.analysers import DEFAULT_FUNCTION_ANALYSERS as DEFAULTS
 
 Plugin = Union[Assertor, CustomFunctionAnalyser]
 
@@ -22,7 +22,7 @@ class Plugins:
     @property
     def custom_function_handler(self) -> CustomFunctionHandler:
         if self._handler is None:
-            self._handler = CustomFunctionHandler(BUILTINS, self.analysers)
+            self._handler = CustomFunctionHandler(DEFAULTS, self.analysers)
 
         return self._handler
 
diff --git a/ratter/plugins/analysers/__init__.py b/ratter/plugins/analysers/__init__.py
index 3dae5b2c..354c9fd1 100644
--- a/ratter/plugins/analysers/__init__.py
+++ b/ratter/plugins/analysers/__init__.py
@@ -5,11 +5,13 @@
     SetattrAnalyser,
     SortedAnalyser,
 )
+from ratter.plugins.analysers.collections import DefaultDictAnalyser
 
-BUILTIN_FUNCTION_ANALYSERS = (
+DEFAULT_FUNCTION_ANALYSERS = (
     GetattrAnalyser(),
     SetattrAnalyser(),
     HasattrAnalyser(),
     DelattrAnalyser(),
     SortedAnalyser(),
+    DefaultDictAnalyser(),
 )
diff --git a/ratter/plugins/analysers/collections.py b/ratter/plugins/analysers/collections.py
new file mode 100644
index 00000000..840144be
--- /dev/null
+++ b/ratter/plugins/analysers/collections.py
@@ -0,0 +1,56 @@
+"""CustomFunctionAnalyser for Python's collections."""
+
+import ast
+
+from ratter.analyser.base import CustomFunctionAnalyser
+from ratter.analyser.context import Context
+from ratter.analyser.types import FuncOrAsyncFunc, FunctionIR
+
+
+class DefaultDictAnalyser(CustomFunctionAnalyser):
+    @property
+    def name(self) -> str:
+        return "defaultdict"
+
+    @property
+    def qualified_name(self) -> str:
+        return "collections.defaultdict"
+
+    def on_def(self, name: str, node: FuncOrAsyncFunc, ctx: Context) -> FunctionIR:
+        return super().on_def(name, node, ctx)
+
+    def on_call(self, name: str, node: ast.Call, ctx: Context) -> FunctionIR:
+        # HACK Avoid circular import
+        from ratter.analyser.function import FunctionAnalyser
+
+        if len(node.args) == 0:
+            return super().on_call(name, node, ctx)
+
+        default_factory = node.args[0]
+
+        # Three cases:
+        #   1. A named callable         ->  Add call to name
+        #   2. A lambda                 ->  Fold-in results of body
+        #   3. Non-lambda expression    ->  Visit
+
+        fn_analyser = FunctionAnalyser(ast.FunctionDef(), ctx)
+
+        if isinstance(default_factory, (ast.Name, ast.Attribute)):
+            target = ast.Call(
+                func=default_factory,
+                args=[],
+                keywords=[],
+            )
+
+            # Populate ast.AST fields
+            target._fields = tuple()
+            target.lineno = default_factory.lineno
+            target.col_offset = default_factory.col_offset
+        elif isinstance(default_factory, ast.Lambda):
+            target = default_factory.body
+        else:
+            target = default_factory
+
+        fn_analyser.visit(target)
+
+        return fn_analyser.func_ir


Test patch:
diff --git a/tests/plugins/analysers/test_collections.py b/tests/plugins/analysers/test_collections.py
new file mode 100644
index 00000000..d3b3d97d
--- /dev/null
+++ b/tests/plugins/analysers/test_collections.py
@@ -0,0 +1,144 @@
+import pytest
+
+from ratter.analyser.context import (
+    RootContext,
+    Builtin,
+    Call,
+    Func,
+    Name,
+)
+from ratter.analyser.file import FileAnalyser
+from ratter.plugins import plugins
+from ratter.plugins.analysers.collections import DefaultDictAnalyser
+
+
+class TestCustomCollectionsAnalysers:
+
+    @pytest.fixture(autouse=True)
+    def apply_plugins(self):
+        plugins.register(DefaultDictAnalyser())
+
+    def test_defaultdict_no_factory(self, parse):
+        _ast = parse("""
+            def a_func(arg):
+                d = defaultdict()
+        """)
+        results = FileAnalyser(_ast, RootContext(_ast)).analyse()
+
+        expected = {
+            Func("a_func", ["arg"], None, None): {
+                "gets": set(),
+                "sets": {
+                    Name("d"),
+                },
+                "dels": set(),
+                "calls": set(),
+            }
+        }
+
+        assert results == expected
+
+    def test_defaultdict_named_factory(self, parse):
+        _ast = parse("""
+            def a_func(arg):
+                d = defaultdict(factory)
+
+            def factory():
+                return hello.results
+        """)
+        results = FileAnalyser(_ast, RootContext(_ast)).analyse()
+
+        a_func = Func("a_func", ["arg"], None, None)
+        factory = Func("factory", [], None, None)
+        expected = {
+            a_func: {
+                "gets": set(),
+                "sets": {
+                    Name("d"),
+                },
+                "dels": set(),
+                "calls": {
+                    Call("factory()", [], {}, target=factory)
+                },
+            },
+            factory: {
+                "gets": {
+                    Name("hello.results", "hello"),
+                },
+                "sets": set(),
+                "dels": set(),
+                "calls": set(),
+            }
+        }
+
+        assert results == expected
+
+    def test_defaultdict_lambda_factory(self, parse):
+        # Lambda to literal
+        _ast = parse("""
+            def a_func(arg):
+                d = defaultdict(lambda: 0)
+        """)
+        results = FileAnalyser(_ast, RootContext(_ast)).analyse()
+
+        expected = {
+            Func("a_func", ["arg"], None, None): {
+                "gets": set(),
+                "sets": {
+                    Name("d"),
+                },
+                "dels": set(),
+                "calls": set(),
+            }
+        }
+
+        print(results)
+        print(expected)
+        assert results == expected
+
+        # Lambda to attr
+        _ast = parse("""
+            def a_func(arg):
+                d = defaultdict(lambda: arg.attr)
+        """)
+        results = FileAnalyser(_ast, RootContext(_ast)).analyse()
+
+        expected = {
+            Func("a_func", ["arg"], None, None): {
+                "gets": {
+                    Name("arg.attr", "arg")
+                },
+                "sets": {
+                    Name("d"),
+                },
+                "dels": set(),
+                "calls": set(),
+            }
+        }
+
+        print(results)
+        print(expected)
+        assert results == expected
+
+    def test_defaultdict_nested_factory(self, parse):
+        _ast = parse("""
+            def a_func(arg):
+                d = defaultdict(defaultdict(list))
+        """)
+        results = FileAnalyser(_ast, RootContext(_ast)).analyse()
+
+        list_builtin = Builtin("list", has_affect=False)
+        expected = {
+            Func("a_func", ["arg"], None, None): {
+                "gets": set(),
+                "sets": {
+                    Name("d"),
+                },
+                "dels": set(),
+                "calls": {
+                    Call("list()", [], {}, target=list_builtin)
+                },
+            }
+        }
+
+        assert results == expected