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
@@ -14,8 +14,12 @@
 from codeop import CommandCompiler, Compile
 import re
 import tokenize
-from typing import List, Tuple, Optional, Any
+from typing import List, Tuple, Optional, Any, TYPE_CHECKING
 import warnings
+import io
+
+if TYPE_CHECKING:
+    from IPython.core.interactiveshell import InteractiveShell
 
 _indent_re = re.compile(r'^[ \t]+')
 
@@ -32,6 +36,7 @@ def leading_empty_lines(lines):
             return lines[i:]
     return lines
 
+
 def leading_indent(lines):
     """Remove leading indentation.
 
@@ -183,6 +188,84 @@ 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)
+
+
+class AutoBalancer:
+    """Callable class, for single line cell, add necessary [{( and )}] to balance."""
+
+    has_side_effects: bool = True
+
+    def __init__(self, shell=None, default=False):
+        self._shell = shell
+        self._default = default
+
+    def __call__(self, lines):
+        """For single line cell, add necessary [{( and )}] to balance."""
+
+        if len(lines) != 1:
+            # apply only for single line cell
+            return lines
+
+        if self._shell is None:
+            if not self._default:
+                return lines
+        elif not self._shell.autobalance:
+            return lines
+
+        modified, new_line = _autobalance_line(lines[0])
+        if modified:
+            if self._shell is not None:
+                self._shell.auto_rewrite_input(new_line)
+            return [new_line]
+        return lines
+
+
 class TokenTransformBase:
     """Base class for transformations which examine tokens.
 
@@ -570,7 +653,10 @@ 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.cleanup_transforms = [
             leading_empty_lines,
             leading_indent,
@@ -579,6 +665,7 @@ def __init__(self):
         ]
         self.line_transforms = [
             cell_magic,
+            AutoBalancer(shell),
         ]
         self.token_transformers = [
             MagicAssign,

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_line.py b/IPython/core/tests/test_inputtransformer2_line.py
@@ -93,6 +93,13 @@ def a():
 """,
 )
 
+AUTOBALANCE = (
+    ("(1+1)/2", "(1+1)/2"),
+    ("1+1)/2", "(1+1)/2"),
+    ("1+1/(2+3", "1+1/(2+3)"),
+    ("i+1)/2 for i in range(10)]", "[(i+1)/2 for i in range(10)]"),
+)
+
 
 def test_ipython_prompt():
     for sample, expected in [
@@ -165,3 +172,15 @@ def test_leading_empty_lines():
 def test_crlf_magic():
     for sample, expected in [CRLF_MAGIC]:
         assert ipt2.cell_magic(sample) == expected
+
+
+def test_autobalance_no_changes():
+    autobalance = ipt2.AutoBalancer(default=False)
+    for sample, _ in AUTOBALANCE:
+        assert autobalance([sample]) == [sample]
+
+
+def test_autobalance_changes():
+    autobalance = ipt2.AutoBalancer(default=True)
+    for sample, expected in AUTOBALANCE:
+        assert autobalance([sample]) == [expected]
diff --git a/IPython/core/tests/test_interactiveshell.py b/IPython/core/tests/test_interactiveshell.py
@@ -791,6 +791,7 @@ class TestMiscTransform(unittest.TestCase):
     def test_transform_only_once(self):
         cleanup = 0
         line_t = 0
+
         def count_cleanup(lines):
             nonlocal cleanup
             cleanup += 1