chore(iast): fix for dict add inplace

benchmarks/bm/iast_fixtures/str_methods.py

@@ -55,6 +55,36 @@ def do_operator_add_inplace_params(a, b):
     return a
 
 
+def do_operator_add_inplace_dict_key(a, key, b):
+    a[key] += b
+    return a[key]
+
+
+def _get_dictionary(dictionary):
+    return dictionary
+
+
+def do_operator_add_inplace_dict_key_from_function(a, key, b):
+    _get_dictionary(a)[key] += b
+    return a[key]
+
+
+class MyClassWithDict(object):
+    data = {}
+
+    def __init__(self, my_cool_dict):
+        self.data = my_cool_dict
+
+    def get_data(self):
+        return self.data
+
+
+def do_operator_add_inplace_dict_key_from_class(a, key, b):
+    my_cool_data = MyClassWithDict(a)
+    my_cool_data.get_data()[key] += b
+    return a[key]
+
+
 def do_operator_add_inplace_3_params(a, b, c):
     a += b
     a += c

ddtrace/appsec/_iast/_ast/visitor.py

@@ -19,7 +19,6 @@
 from ..constants import DEFAULT_WEAK_RANDOMNESS_FUNCTIONS
 
 
-PY3 = sys.version_info[0] >= 3
 PY39_PLUS = sys.version_info >= (3, 9, 0)
 
 _PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX
@@ -230,7 +229,7 @@ def _merge_dicts(*args_functions: Set[str]) -> Set[str]:
 
     @staticmethod
     def _is_string_node(node: Any) -> bool:
-        if PY3 and (isinstance(node, ast.Constant) and isinstance(node.value, IAST.TEXT_TYPES)):
+        if isinstance(node, ast.Constant) and isinstance(node.value, IAST.TEXT_TYPES):
             return True
 
         return False
@@ -652,46 +651,39 @@ def visit_AugAssign(self, augassign_node: ast.AugAssign) -> Any:
         """
         Replace an inplace add or multiply (+= / *=)
         """
-        if isinstance(augassign_node.target, ast.Subscript):
-            # Can't augassign to function call, ignore this node
-            augassign_node.target.avoid_convert = True  # type: ignore[attr-defined]
-            self.generic_visit(augassign_node)
-            return augassign_node
-
         self.generic_visit(augassign_node)
+
         if augassign_node.op.__class__ == ast.Add:
             # Optimization: ignore augassigns where the right side term is an integer since
             # they can't apply to strings
             if self._is_numeric_node(augassign_node.value):
                 return augassign_node
 
             replacement_func = self._aspect_operators["INPLACE_ADD"]
-        else:
-            return augassign_node
-
-        # We must change the ctx of the target and value to Load() for using
-        # them as function arguments while keeping it as Store() for the
-        # Assign.targets, thus the manual copy
 
-        func_arg1 = copy.deepcopy(augassign_node.target)
-        func_arg1.ctx = ast.Load()
-        func_arg2 = copy.deepcopy(augassign_node.value)
-        func_arg2.ctx = ast.Load()  # type: ignore[attr-defined]
+            # Regular inplace add for non-subscript targets
+            func_arg1 = copy.deepcopy(augassign_node.target)
+            func_arg1.ctx = ast.Load()
+            func_arg2 = copy.deepcopy(augassign_node.value)
+            if hasattr(func_arg2, "ctx"):
+                func_arg2.ctx = ast.Load()
+
+            call_node = self._call_node(
+                augassign_node,
+                func=self._attr_node(augassign_node, replacement_func),
+                args=[func_arg1, func_arg2],
+            )
 
-        call_node = self._call_node(
-            augassign_node,
-            func=self._attr_node(augassign_node, replacement_func),
-            args=[func_arg1, func_arg2],
-        )
+            self.ast_modified = True
+            return self._node(
+                ast.Assign,
+                augassign_node,
+                targets=[augassign_node.target],
+                value=call_node,
+                type_comment=None,
+            )
 
-        self.ast_modified = True
-        return self._node(
-            ast.Assign,
-            augassign_node,
-            targets=[augassign_node.target],
-            value=call_node,
-            type_comment=None,
-        )
+        return augassign_node
 
     def visit_FormattedValue(self, fmt_value_node: ast.FormattedValue) -> Any:
         """