ipython · jb-leger · May 19, 2023 · May 19, 2023 · May 19, 2023 · May 19, 2023
diff --git a/IPython/core/inputtransformer2.py b/IPython/core/inputtransformer2.py
@@ -16,10 +16,16 @@
 import tokenize
 from typing import List, Tuple, Optional, Any
 import warnings
+import io
+
+# only for type checking, do not import this, it lead to a circular
+# import
+if False:
+    from IPython.core.interactiveshell import InteractiveShell
 
 _indent_re = re.compile(r'^[ \t]+')
 
-def leading_empty_lines(lines):
+def leading_empty_lines(lines, **kwargs):
     """Remove leading empty lines
 
     If the leading lines are empty or contain only whitespace, they will be
@@ -32,7 +38,8 @@ def leading_empty_lines(lines):
             return lines[i:]
     return lines
 
-def leading_indent(lines):
+
+def leading_indent(lines, **kwargs):
     """Remove leading indentation.
 
     If the first line starts with a spaces or tabs, the same whitespace will be
@@ -77,7 +84,7 @@ def __init__(self, prompt_re, initial_re=None):
     def _strip(self, lines):
         return [self.prompt_re.sub('', l, count=1) for l in lines]
 
-    def __call__(self, lines):
+    def __call__(self, lines, **kwargs):
         if not lines:
             return lines
         if self.initial_re.match(lines[0]) or \
@@ -115,7 +122,7 @@ def __call__(self, lines):
 )
 
 
-def cell_magic(lines):
+def cell_magic(lines, **kwargs):
     if not lines or not lines[0].startswith('%%'):
         return lines
     if re.match(r'%%\w+\?', lines[0]):
@@ -183,6 +190,74 @@ def assemble_continued_line(lines, start: Tuple[int, int], end_line: int):
     return ' '.join([p.rstrip()[:-1] for p in parts[:-1]]  # Strip backslash+newline
                     + [parts[-1].rstrip()])         # Strip newline from last line
 
+_AUTOBALANCE_MARKERS = (("(", ")"), ("{", "}"), ("[", "]"))
+_AUTOBALANCE_OPEN = {x: i for i, (x, _) in enumerate(_AUTOBALANCE_MARKERS)}
+_AUTOBALANCE_CLOSE = {x: i for i, (_, x) in enumerate(_AUTOBALANCE_MARKERS)}
+
+
+def _autobalance_line(inpt_code):
+    """Add necessary [{( in the begin of the expr and )}] at the end to balance."""
+
+    tokens = tokenize.generate_tokens(io.StringIO(inpt_code).readline)
+    closed_without_open = []
+    opened_without_close = []
+    begin_expr = 0
+    while True:
+        try:
+            token = next(tokens)
+        except (StopIteration, tokenize.TokenError):
+            break
+        if token.type == tokenize.OP:
+            if token.string in _AUTOBALANCE_OPEN:
+                opened_without_close.append(_AUTOBALANCE_OPEN[token.string])
+            elif token.string in _AUTOBALANCE_CLOSE:
+                if opened_without_close:
+                    last_opened = opened_without_close.pop(-1)
+                    if last_opened != _AUTOBALANCE_CLOSE[token.string]:
+                        # can not be balanced only adding in the begin and end
+                        # of expr
+                        return (False, inpt_code)
+                else:
+                    closed_without_open.append(_AUTOBALANCE_CLOSE[token.string])
+            elif token.string == "=":
+                if not opened_without_close:
+                    # this is a assigment
+                    begin_expr = token.end[1]
+
+    if not opened_without_close and not closed_without_open:
+        # no needed changes
+        return (False, inpt_code)
+
+    new_code = (
+        inpt_code[:begin_expr]
+        + (" " if begin_expr else "")
+        + "".join(_AUTOBALANCE_MARKERS[k][0] for k in closed_without_open[::-1])
+        + inpt_code[begin_expr:].strip()
+        + "".join(_AUTOBALANCE_MARKERS[k][1] for k in opened_without_close)
+    )
+    return (True, new_code)
+
+
+def autobalance(lines, **kwargs):
+    """For single line cell, add necessary [{( and )}] to balance."""
+
+    shell = kwargs.get("shell")
+    display = kwargs.get("display", False)
+
+    if len(lines) != 1:
+        # apply only for single line cell
+        return lines
+
+    if shell is None or not shell.autobalance:
+        return lines
+
+    modified, new_line = _autobalance_line(lines[0])
+    if modified and display:
+        shell.auto_rewrite_input(new_line)
+        return [new_line]
+    return lines
+
+
 class TokenTransformBase:
     """Base class for transformations which examine tokens.
 
@@ -248,7 +323,7 @@ def find(cls, tokens_by_line):
                     and (line[assign_ix+2].type == tokenize.NAME):
                 return cls(line[assign_ix+1].start)
 
-    def transform(self, lines: List[str]):
+    def transform(self, lines: List[str], **kwargs):
         """Transform a magic assignment found by the ``find()`` classmethod.
         """
         start_line, start_col = self.start_line, self.start_col
@@ -287,7 +362,7 @@ def find(cls, tokens_by_line):
                         break
                     ix += 1
 
-    def transform(self, lines: List[str]):
+    def transform(self, lines: List[str], **kwargs):
         """Transform a system assignment found by the ``find()`` classmethod.
         """
         start_line, start_col = self.start_line, self.start_col
@@ -404,7 +479,7 @@ def find(cls, tokens_by_line):
             if line[ix].string in ESCAPE_SINGLES:
                 return cls(line[ix].start)
 
-    def transform(self, lines):
+    def transform(self, lines, **kwargs):
         """Transform an escaped line found by the ``find()`` classmethod.
         """
         start_line, start_col = self.start_line, self.start_col
@@ -465,7 +540,7 @@ def find(cls, tokens_by_line):
                     ix += 1
                 return cls(line[ix].start, line[-2].start)
 
-    def transform(self, lines):
+    def transform(self, lines, **kwargs):
         """Transform a help command found by the ``find()`` classmethod.
         """
 
@@ -570,7 +645,11 @@ class TransformerManager:
     The key methods for external use are ``transform_cell()``
     and ``check_complete()``.
     """
-    def __init__(self):
+
+    shell: Optional["InteractiveShell"]
+
+    def __init__(self, shell=None):
+        self.shell = shell
         self.cleanup_transforms = [
             leading_empty_lines,
             leading_indent,
@@ -579,6 +658,7 @@ def __init__(self):
         ]
         self.line_transforms = [
             cell_magic,
+            autobalance,
         ]
         self.token_transformers = [
             MagicAssign,
@@ -587,7 +667,7 @@ def __init__(self):
             HelpEnd,
         ]
 
-    def do_one_token_transform(self, lines):
+    def do_one_token_transform(self, lines, display=False):
         """Find and run the transform earliest in the code.
 
         Returns (changed, lines).
@@ -614,14 +694,16 @@ def do_one_token_transform(self, lines):
         ordered_transformers = sorted(candidates, key=TokenTransformBase.sortby)
         for transformer in ordered_transformers:
             try:
-                return True, transformer.transform(lines)
+                return True, transformer.transform(
+                    lines, shell=self.shell, display=display
+                )
             except SyntaxError:
                 pass
         return False, lines
 
-    def do_token_transforms(self, lines):
+    def do_token_transforms(self, lines, display=False):
         for _ in range(TRANSFORM_LOOP_LIMIT):
-            changed, lines = self.do_one_token_transform(lines)
+            changed, lines = self.do_one_token_transform(lines, display=display)
             if not changed:
                 return lines
 
@@ -634,7 +716,7 @@ def transform_cell(self, cell: str) -> str:
             cell += '\n'  # Ensure the cell has a trailing newline
         lines = cell.splitlines(keepends=True)
         for transform in self.cleanup_transforms + self.line_transforms:
-            lines = transform(lines)
+            lines = transform(lines, shell=self.shell, display=True)
 
         lines = self.do_token_transforms(lines)
         return ''.join(lines)
@@ -684,7 +766,7 @@ def check_complete(self, cell: str):
         try:
             for transform in self.cleanup_transforms:
                 if not getattr(transform, 'has_side_effects', False):
-                    lines = transform(lines)
+                    lines = transform(lines, shell=self.shell, display=False)
         except SyntaxError:
             return 'invalid', None
 
@@ -698,8 +780,8 @@ def check_complete(self, cell: str):
         try:
             for transform in self.line_transforms:
                 if not getattr(transform, 'has_side_effects', False):
-                    lines = transform(lines)
-            lines = self.do_token_transforms(lines)
+                    lines = transform(lines, shell=self.shell, display=False)
+            lines = self.do_token_transforms(lines, display=False)
         except SyntaxError:
             return 'invalid', None
 

diff --git a/IPython/core/interactiveshell.py b/IPython/core/interactiveshell.py
@@ -71,7 +71,7 @@
 from IPython.core.extensions import ExtensionManager
 from IPython.core.formatters import DisplayFormatter
 from IPython.core.history import HistoryManager
-from IPython.core.inputtransformer2 import ESC_MAGIC, ESC_MAGIC2
+from IPython.core.inputtransformer2 import ESC_MAGIC, ESC_MAGIC2, TransformerManager
 from IPython.core.logger import Logger
 from IPython.core.macro import Macro
 from IPython.core.payload import PayloadManager
@@ -306,6 +306,16 @@ class InteractiveShell(SingletonConfigurable):
         """
     ).tag(config=True)
 
+    autobalance = Bool(
+        False,
+        help="""
+        Add automatically opening and closing parenthesis, braces and brackets
+        at the begin matching respectively closing and opening ones. Opening
+        symbols are added in the begin of the expression, and closings symbols
+        at the end of the line. E.g. `1+2)/2]*(2+3` becomes `[(1+2)/2]*(2+3)`.
+        """,
+    ).tag(config=True)
+
     autoindent = Bool(True, help=
         """
         Autoindent IPython code entered interactively.
@@ -430,8 +440,9 @@ def _exiter_default(self):
     ipython_dir= Unicode('').tag(config=True) # Set to get_ipython_dir() in __init__
 
     # Used to transform cells before running them, and check whether code is complete
-    input_transformer_manager = Instance('IPython.core.inputtransformer2.TransformerManager',
-                                         ())
+    input_transformer_manager = Instance(
+        "IPython.core.inputtransformer2.TransformerManager", allow_none=True
+    )
 
     @property
     def input_transformers_cleanup(self):
@@ -484,9 +495,8 @@ def input_splitter(self):
         will be displayed as regular output instead."""
     ).tag(config=True)
 
-
-    show_rewritten_input = Bool(True,
-        help="Show rewritten input, e.g. for autocall."
+    show_rewritten_input = Bool(
+        True, help="Show rewritten input, e.g. for autocall and autobalance."
     ).tag(config=True)
 
     quiet = Bool(False).tag(config=True)
@@ -593,6 +603,7 @@ def __init__(self, ipython_dir=None, profile_dir=None,
         self.init_history()
         self.init_encoding()
         self.init_prefilter()
+        self.init_input_transformer_manager()
 
         self.init_syntax_highlighting()
         self.init_hooks()
@@ -2740,6 +2751,13 @@ def auto_rewrite_input(self, cmd):
         print("------> " + cmd)
 
     #-------------------------------------------------------------------------
+    # Things related to input_transformer_manager
+    # -------------------------------------------------------------------------
+
+    def init_input_transformer_manager(self):
+        self.input_transformer_manager = TransformerManager(shell=self)
+
+    # -------------------------------------------------------------------------
     # Things related to extracting values/expressions from kernel and user_ns
     #-------------------------------------------------------------------------
 

diff --git a/IPython/core/magics/auto.py b/IPython/core/magics/auto.py
@@ -142,3 +142,32 @@ def errorMessage() -> str:
                     self.shell.autocall = self._magic_state.autocall_save = 1
 
         print("Automatic calling is:", list(valid_modes.values())[self.shell.autocall])
+
+    @line_magic
+    def autobalance(self, parameter_s=""):
+        """Automatically balance parenthesis, braces and brackets.
+
+        On single line cell, this mode add automatically needed symbol at the
+        begin of the expression and the end of the line. If symbols need to
+        be added at other place, there exists a ambiguity, and this
+        transformation is not applied.
+
+        Without arguments toggles on/off. With arguments it sets the value, and
+        you can use any of (case insensitive):
+
+         - on, 1, True: to activate
+
+         - off, 0, False: to deactivate.
+
+        When symbols are added, the rewritten line is displayed. See
+        shell configuration `show_rewritten_input` to change this behavior.
+        """
+
+        arg = parameter_s.lower()
+        if arg in ("on", "1", "true"):
+            self.shell.autobalance = True
+        elif arg in ("off", "0", "false"):
+            self.shell.autobalance = False
+        else:
+            self.shell.autobalance = not self.shell.autobalance
+        print("Autobalance " + ("enabled" if self.shell.autobalance else "disabled"))
diff --git a/IPython/core/tests/test_inputtransformer2.py b/IPython/core/tests/test_inputtransformer2.py
@@ -160,7 +160,7 @@ def test():
 )
 
 
-def null_cleanup_transformer(lines):
+def null_cleanup_transformer(lines, **kwargs):
     """
     A cleanup transform that returns an empty list.
     """

diff --git a/IPython/core/tests/test_interactiveshell.py b/IPython/core/tests/test_interactiveshell.py
@@ -791,12 +791,13 @@ class TestMiscTransform(unittest.TestCase):
     def test_transform_only_once(self):
         cleanup = 0
         line_t = 0
-        def count_cleanup(lines):
+
+        def count_cleanup(lines, **kwargs):
             nonlocal cleanup
             cleanup += 1
             return lines
 
-        def count_line_t(lines):
+        def count_line_t(lines, **kwargs):
             nonlocal line_t
             line_t += 1
             return lines
-Original file line number
+Diff line change
@@ Expand Up / @@ -160,7 +160,7 @@ def test(): @@
     )
-    def null_cleanup_transformer(lines):
+    def null_cleanup_transformer(lines, **kwargs):
         """
         A cleanup transform that returns an empty list.
         """
@@ Expand Down @@