Repository: mpmath/mpmath (python)

Issue:
mpf_bernoulli sometimes results sometimes non-normalized 
For example:
```pycon
>>> from mpmath import libmp
>>> libmp.mpf_bernoulli(10, 336, rnd='d')
(0, mpz(5970000938847648365496824257597517243193899889464461616355802030946321480529880395668828220437108430425246665198001), -385, 382)
>>> libmp.mpf_bernoulli(10, 336, rnd='u')
(0, mpz(84838815991583492824145540930622889322789602167979149049242011400143631484097005506851420199682087719), -339, 336)
```
while
```pycon
>>> from mpmath import libmp
>>> libmp.mpf_bernoulli(10, 336, rnd='u')
(0, mpz(5970000938847648365496824257597517243193899889464461616355802030946321480529880395668828220437108430425246665198001), -385, 382)
```

It seems that main loop assumes default `rnd=round_fast` value:
https://github.com/mpmath/mpmath/blob/9ea52503dd085ff8a15e17b1fc4f2d0cc36c2bb1/mpmath/libmp/gammazeta.py#L418-L450

Pull request:
Misc fixes
* ensure mpf_bernoulli() returns normalized answer, closes #928
* fix source transformation in IPython, add tests
* use ast transformation in wrap_float_literals()

Patch:
diff --git a/mpmath/__main__.py b/mpmath/__main__.py
index c0d5b46..273767f 100644
--- a/mpmath/__main__.py
+++ b/mpmath/__main__.py
@@ -81,10 +81,13 @@ def main():
         for l in lines:
             shell.run_cell(l, silent=True)
         if not args.no_wrap_floats:
-            shell.run_cell('from mpmath._interactive import wrap_float_literals')
-            shell.run_cell('ip = get_ipython()')
-            shell.run_cell('ip.input_transformers_post.append(wrap_float_literals)')
-            shell.run_cell('del ip')
+            source = """
+from mpmath._interactive import wrap_float_literals
+ip = get_ipython()
+ip.input_transformers_post.append(wrap_float_literals)
+del ip
+"""
+            shell.run_cell(source)
         app.start()
     else:
         ast_transformers = []
@@ -116,8 +119,15 @@ def __init__(self, ast_transformers=[],
                 self.source_transformers = source_transformers
 
             def runsource(self, source, filename='<input>', symbol='single'):
-                for t in self.source_transformers:
-                    source = '\n'.join(t(source.splitlines()))
+                if self.source_transformers:
+                    last_line = source.endswith("\n")  # signals the end of a block
+                    try:
+                        for t in self.source_transformers:
+                            source = ''.join(t(source.splitlines(keepends=True)))
+                    except SyntaxError:
+                        pass  # XXX: emit warning?
+                    if last_line:
+                        source += "\n"
 
                 try:
                     code = self.compile(source, filename, symbol)
@@ -137,8 +147,8 @@ def runsource(self, source, filename='<input>', symbol='single'):
                         tree = t.visit(tree)
                     ast.fix_missing_locations(tree)
                     source = ast.unparse(tree)
+                    source += "\n"
 
-                source += "\n"
                 return super().runsource(source, filename=filename, symbol=symbol)
 
         c = MpmathConsole(ast_transformers=ast_transformers,
diff --git a/mpmath/_interactive.py b/mpmath/_interactive.py
index 2df5cff..09d40e6 100644
--- a/mpmath/_interactive.py
+++ b/mpmath/_interactive.py
@@ -27,16 +27,34 @@ def is_integer(x):
         return self.generic_visit(node)
 
 
+class _WrapFloats(ast.NodeTransformer):
+    """Wrap float literals by calls to specified type."""
+    def __init__(self, lines, type):
+        super().__init__()
+        self.lines = lines
+        self.type = type
+
+    def visit_Constant(self, node):
+        if isinstance(node.value, (float, complex)):
+            line = self.lines[node.lineno - 1]
+            value = line[node.col_offset:node.end_col_offset]
+            is_complex = value.endswith(('j', 'J'))
+            if is_complex:
+                value = value[:-1]
+            value = ast.Constant(value)
+            value = ast.Call(ast.Name(self.type, ast.Load()), [value], [])
+            if is_complex:
+                value = ast.BinOp(left=value, op=ast.Mult(),
+                                  right=ast.Constant(1j))
+            return value
+        return node
+
+
 def wrap_float_literals(lines):
     """Wraps all float/complex literals with mpmath classes."""
-    result = []
-    source = "\n".join(lines)
-    g = tokenize.tokenize(io.BytesIO(source.encode()).readline)
-    for toknum, tokval, *_ in g:
-        if toknum == tokenize.NUMBER:
-            if any(_ in tokval for _ in ['j', 'J']):
-                tokval = f"mpc(imag=mpf('{tokval[:-1]}'))"
-            elif any(_ in tokval for _ in ['.', 'e', 'E']):
-                tokval = f"mpf('{tokval}')"
-        result.append((toknum, tokval, *_))
-    return tokenize.untokenize(result).decode().splitlines()
+    source = ''.join(lines)
+    tree = ast.parse(source)
+    tree = _WrapFloats(lines, 'mpf').visit(tree)
+    ast.fix_missing_locations(tree)
+    source = ast.unparse(tree)
+    return source.splitlines(keepends=True)
diff --git a/mpmath/libmp/gammazeta.py b/mpmath/libmp/gammazeta.py
index 2cb2721..ba22020 100644
--- a/mpmath/libmp/gammazeta.py
+++ b/mpmath/libmp/gammazeta.py
@@ -448,7 +448,7 @@ def mpf_bernoulli(n, prec, rnd=round_fast, plus=False):
         if m > 6:
             bin1 = bin1 * ((2+m)*(3+m)) // ((m-7)*(m-6))
         state[:] = [m, bin, bin1]
-    return numbers[n]
+    return mpf_pos(numbers[n], prec, rnd)
 
 def mpf_bernoulli_huge(n, prec, rnd=round_fast):
     wp = prec + 10


Test patch:
diff --git a/mpmath/tests/test_cli.py b/mpmath/tests/test_cli.py
index 9a7b430..f643369 100644
--- a/mpmath/tests/test_cli.py
+++ b/mpmath/tests/test_cli.py
@@ -90,6 +90,10 @@ def test_ipython_console_wrap_floats():
     assert c.expect_exact('\r\nIn [1]: ') == 0
     assert c.send('10.9\r\n') == 6
     assert c.expect_exact("\r\nOut[1]: mpf('10.899999999999999999999999999995')\r\n\r\nIn [2]: ") == 0
+    assert c.send('def f():\r\n    x = 1.1\n    return x + 1\n\r\n\n') == 42
+    assert c.expect_exact("\r\n\r\nIn [3]: ") == 0
+    assert c.send('f()\r\n') == 5
+    assert c.expect_exact("\r\nOut[3]: mpf('2.0999999999999999999999999999987')\r\n\r\nIn [4]: ") == 0
 
 
 def test_bare_console_wrap_floats():
@@ -120,6 +124,10 @@ def test_bare_console_pretty():
     assert c.expect_exact("10.899999999999999999999999999995\r\n>>> ") == 0
     assert c.send("def f():\r\n  x = ?\r\n\r\n") == 21
     assert c.expect('SyntaxError:') == 0
+    assert c.send('def f():\r\n    x = 1.1\n    return x + 1\n\r\n\n') == 42
+    assert c.expect_exact('>>> ') == 0
+    assert c.send("f()\r\n") == 5
+    assert c.expect_exact('2.0999999999999999999999999999987\r\n>>> ') == 0
 
 
 def test_mpmath_version():
diff --git a/mpmath/tests/test_gammazeta.py b/mpmath/tests/test_gammazeta.py
index cb1acc4..8bad071 100644
--- a/mpmath/tests/test_gammazeta.py
+++ b/mpmath/tests/test_gammazeta.py
@@ -41,6 +41,7 @@ def test_bernoulli(plus):
     assert bernoulli(5, plus) == 0
     assert bernoulli(6, plus).ae(1./42)
     assert str(bernoulli(10, plus)) == '0.0757575757575758'
+    assert repr(bernoulli(10, plus)) == "mpf('0.07575757575757576')"
     assert str(bernoulli(234, plus)) == '7.62772793964344e+267'
     assert str(bernoulli(10**5, plus)) == '-5.82229431461335e+376755'
     assert str(bernoulli(10**8+2, plus)) == '1.19570355039953e+676752584'