Add cli option to wrap integer literals with Integer

diofant/__main__.py

@@ -17,7 +17,7 @@
 
 from diofant import __version__
 from diofant.interactive.session import (AutomaticSymbols,
-                                         IntegerDivisionWrapper,
+                                         IntegerDivisionWrapper, WrapInts,
                                          unicode_identifiers,
                                          wrap_float_literals)
 
@@ -40,6 +40,8 @@
                     action='store_true')
 parser.add_argument('--wrap-floats', help='Wrap float literals with Float',
                     action='store_true')
+parser.add_argument('--wrap-ints', help='Wrap integer literals with Integer',
+                    action='store_true')
 parser.add_argument('-V', '--version',
                     help='Print the Diofant version and exit',
                     action='store_true')
@@ -59,6 +61,9 @@ def main():
              "f, g, h = symbols('f g h', cls=Function)",
              'init_printing(pretty_print=True, use_unicode=True)']
 
+    if args.wrap_ints:
+        args.no_wrap_division = True
+
     try:
         import IPython
         import traitlets
@@ -95,6 +100,11 @@ def main():
             shell.run_cell('ip = get_ipython()')
             shell.run_cell('ip.input_transformers_cleanup.append(wrap_float_literals)')
             shell.run_cell('del ip')
+        if args.wrap_ints:
+            shell.run_cell('from diofant.interactive.session import WrapInts')
+            shell.run_cell('ip = get_ipython()')
+            shell.run_cell('ip.ast_transformers.append(WrapInts())')
+            shell.run_cell('del ip')
         app.start()
     else:
         ast_transformers = []
@@ -109,6 +119,8 @@ def main():
             source_transformers.append(unicode_identifiers)
         if args.wrap_floats:
             source_transformers.append(wrap_float_literals)
+        if args.wrap_ints:
+            ast_transformers.append(WrapInts())
 
         class DiofantConsole(code.InteractiveConsole):
             """An interactive console with readline support."""

diofant/interactive/session.py

@@ -80,6 +80,16 @@ def visit_Call(self, node):
         return self.generic_visit(node)
 
 
+class WrapInts(ast.NodeTransformer):
+    """Wraps all ints in a call to :class:`~diofant.core.numbers.Integer`."""
+
+    def visit_Constant(self, node):
+        if isinstance(node.value, int) and not isinstance(node.value, bool):
+            return ast.Call(ast.Name('Integer', ast.Load()),
+                            [ast.Constant(repr(node.value))], [])
+        return node
+
+
 _NAMES_MAP = {}
 
 

diofant/tests/test_cli.py

@@ -157,6 +157,26 @@ def test_bare_console_wrap_floats():
     assert c.expect_exact('2.00000000000000\r\n>>> ') == 0
 
 
+def test_ipython_console_wrap_ints():
+    pytest.importorskip('IPython')
+
+    c = Console('python -m diofant --simple-prompt '
+                "--wrap-ints --colors 'NoColor'")
+
+    assert c.expect_exact('\r\nIn [1]: ') == 0
+    assert c.send('repr(10)\r\n') == 10
+    assert c.expect_exact("\r\nOut[1]: \'Integer(10)\'\r\n\r\nIn [2]: ") == 0
+
+
+def test_bare_console_wrap_ints():
+    c = Console('python -m diofant --simple-prompt --no-ipython '
+                "--wrap-ints --colors 'NoColor'")
+
+    assert c.expect_exact('>>> ') == 0
+    assert c.send('repr(10)\r\n') == 10
+    assert c.expect_exact("\'Integer(10)\'\r\n>>> ") == 0
+
+
 def test_diofant_version():
     c = Console('python -m diofant --version')