improve-code-strip-logic

Author: Darshan808

IPython/core/completer.py

@@ -1134,20 +1134,60 @@ def attr_matches(self, text):
     # we simple attribute matching with normal identifiers.
     _ATTR_MATCH_RE = re.compile(r"(.+)\.(\w*)$")
 
+    def _strip_code_before_operator(self, code: str) -> str:
+        operators = ["=", "==", "!=", ">=", "<=", ">", "<"]
+        o_parens = {"(", "[", "{"}
+        c_parens = {")", "]", "}"}
+
+        # Dry-run tokenize to catch errors
+        try:
+            _ = list(tokenize.generate_tokens(iter(code.splitlines()).__next__))
+        except tokenize.TokenError:
+            # Try trimming the expression and retrying
+            trimmed_code = self._trim_expr(code)
+            try:
+                _ = list(
+                    tokenize.generate_tokens(iter(trimmed_code.splitlines()).__next__)
+                )
+                code = trimmed_code
+            except tokenize.TokenError:
+                return code
+
+        tokens = _parse_tokens(code)
+        encountered_operator = False
+        after_operator = []
+        nesting_level = 0
+
+        for t in tokens:
+            if t.type == tokenize.OP:
+                if t.string in o_parens:
+                    nesting_level += 1
+                elif t.string in c_parens:
+                    nesting_level -= 1
+                elif t.string in operators and nesting_level == 0:
+                    encountered_operator = True
+                    after_operator = []
+                    continue
+
+            if encountered_operator:
+                after_operator.append(t.string)
+
+        if encountered_operator:
+            return "".join(after_operator)
+        else:
+            return code
+
     def _attr_matches(
         self, text: str, include_prefix: bool = True
     ) -> tuple[Sequence[str], str]:
         m2 = self._ATTR_MATCH_RE.match(text)
         if not m2:
             return [], ""
         expr, attr = m2.group(1, 2)
-
-        operators = ["=", "==", "!=", ">=", "<=", ">", "<"]
-        # Split by the operator and take the right side
-        if not (expr.startswith("(") and expr.endswith(")")):
-            for op in operators:
-                if op in expr:
-                    expr = expr.split(op)[-1]
+        try:
+            expr = self._strip_code_before_operator(expr)
+        except tokenize.TokenError:
+            pass
 
         obj = self._evaluate_expr(expr)
 

tests/test_completer.py

@@ -662,14 +662,27 @@ def _(line, cursor_pos, expect, message, completion):
                 "Should have completed on `x.upper() == y.`: %s",
                 Completion(15, 15, "upper"),
             )
+            _(
+                "(x.upper() == y.",
+                16,
+                ".upper",
+                "Should have completed on `x.upper() == y.`: %s",
+                Completion(16, 16, "upper"),
+            )
             _(
                 "(x.upper() == y).",
                 17,
                 ".bit_length",
                 "Should have completed on `(x.upper() == y).`: %s",
                 Completion(17, 17, "bit_length"),
             )
-
+            _(
+                "{'==', 'abc'}.",
+                14,
+                ".add",
+                "Should have completed on `{'==', 'abc'}.`: %s",
+                Completion(14, 14, "add"),
+            )
     def test_omit__names(self):
         # also happens to test IPCompleter as a configurable
         ip = get_ipython()