Support binary/octal/hexadecimal string output

Adapted from #414.

Partial fix for #345

Co-authored-by: Sergey B Kirpichev <[email protected]>

* this version doesn't touch argument names (dps->ps)
* no support for float.hex()-style binary exponents

Author: warnerjon12

mpmath/libmp/libmpf.py

@@ -15,7 +15,7 @@
                       MPZ_TWO, MPZ_ZERO, STRICT, gmpy, sage, sage_utils)
 from .libintmath import (bctable, bin_to_radix, giant_steps, isqrt, isqrt_fast,
                          lshift, numeral, rshift, sqrt_fixed, sqrtrem,
-                         trailing, trailtable)
+                         stddigits, trailing, trailtable)
 
 
 class ComplexResult(ValueError):
@@ -1028,13 +1028,13 @@ def mpf_perturb(x, eps_sign, prec, rnd):
 #                              Radix conversion                              #
 #----------------------------------------------------------------------------#
 
-def to_digits_exp(s, dps):
+def to_digits_exp(s, dps, base=10):
     """Helper function for representing the floating-point number s as
-    a decimal with dps digits. Returns (sign, string, exponent) where
-    sign is '' or '-', string is the digit string, and exponent is
-    the decimal exponent as an int.
+    a string with dps digits. Returns (sign, string, exponent) where
+    sign is '' or '-', string is the digit string in the given base,
+    and exponent is the exponent as an int.
 
-    If inexact, the decimal representation is rounded toward zero."""
+    If inexact, the string representation is rounded toward zero."""
 
     # Extract sign first so it doesn't mess up the string digit count
     if s[0]:
@@ -1047,12 +1047,19 @@ def to_digits_exp(s, dps):
     if not man:
         return '', '0', 0
 
-    bitprec = int(dps * math.log(10,2)) + 10
+    if base == 10:
+        blog2 = 3.3219280948873626
+    elif pow(2, blog2 := int(math.log2(base))) == base:
+        pass
+    else:
+        raise NotImplementedError
+
+    bitprec = int(dps * blog2) + 10
 
     # Cut down to size
     # TODO: account for precision when doing this
     exp_from_1 = exp + bc
-    if abs(exp_from_1) > 3500:
+    if base == 10 and abs(exp_from_1) > 3500:
         from .libelefun import mpf_ln2, mpf_ln10
 
         # Set b = int(exp * log(2)/log(10))
@@ -1073,19 +1080,19 @@ def to_digits_exp(s, dps):
     # fixed-point number and then converting that number to
     # a decimal fixed-point number.
     fixprec = max(bitprec - exp - bc, 0)
-    fixdps = int(fixprec / math.log(10,2) + 0.5)
+    fixdps = int(fixprec / blog2 + 0.5)
     sf = to_fixed(s, fixprec)
-    sd = bin_to_radix(sf, fixprec, 10, fixdps)
-    digits = numeral(sd, base=10, size=dps)
+    sb = bin_to_radix(sf, fixprec, base, fixdps)
+    digits = numeral(sb, base=base, size=dps)
 
     exponent += len(digits) - fixdps - 1
     return sign, digits, exponent
 
 def to_str(s, dps, strip_zeros=True, min_fixed=None, max_fixed=None,
-    show_zero_exponent=False):
+    show_zero_exponent=False, base=10):
     """
-    Convert a raw mpf to a decimal floating-point literal with at
-    most `dps` decimal digits in the mantissa (not counting extra zeros
+    Convert a raw mpf to a floating-point literal in the given base
+    with at most `dps` digits in the mantissa (not counting extra zeros
     that may be inserted for visual purposes).
 
     The number will be printed in fixed-point format if the position
@@ -1100,13 +1107,15 @@ def to_str(s, dps, strip_zeros=True, min_fixed=None, max_fixed=None,
     by from_str, float() or Decimal().
     """
 
+    sep = '@' if base > 10 else 'e'
+
     # Special numbers
     if not s[1]:
         if s == fzero:
             if dps: t = '0.0'
             else:   t = '.0'
             if show_zero_exponent:
-                t += 'e+0'
+                t += sep + '+0'
             return t
         if s == finf: return '+inf'
         if s == fninf: return '-inf'
@@ -1118,23 +1127,27 @@ def to_str(s, dps, strip_zeros=True, min_fixed=None, max_fixed=None,
 
     # to_digits_exp rounds to floor.
     # This sometimes kills some instances of "...00001"
-    sign, digits, exponent = to_digits_exp(s, dps+3)
+    sign, digits, exponent = to_digits_exp(s, dps+3, base)
+
+    rnd_digs = stddigits[(base//2 + base%2):base]
 
     # No digits: show only .0; round exponent to nearest
     if not dps:
-        if digits[0] in '56789':
+        if digits[0] in rnd_digs:
             exponent += 1
         digits = ".0"
 
     else:
         # Rounding up kills some instances of "...99999"
-        if len(digits) > dps and digits[dps] in '56789':
+        if len(digits) > dps and digits[dps] in rnd_digs:
             digits = digits[:dps]
             i = dps - 1
-            while i >= 0 and digits[i] == '9':
+            dig = stddigits[base-1]
+            while i >= 0 and digits[i] == dig:
                 i -= 1
             if i >= 0:
-                digits = digits[:i] + str(int(digits[i]) + 1) + '0' * (dps - i - 1)
+                digits = digits[:i] + stddigits[int(digits[i], base) + 1] + \
+                    '0' * (dps - i - 1)
             else:
                 digits = '1' + '0' * (dps - 1)
                 exponent += 1
@@ -1162,9 +1175,19 @@ def to_str(s, dps, strip_zeros=True, min_fixed=None, max_fixed=None,
             if digits[-1] == ".":
                 digits += "0"
 
+    if base == 2:
+        prefix = "0b"
+    elif base == 8:
+        prefix = "0o"
+    elif base == 16:
+        prefix = "0x"
+    else:
+        prefix = ""
+
+    sign += prefix
+
     if exponent == 0 and dps and not show_zero_exponent: return sign + digits
-    if exponent >= 0: return sign + digits + "e+" + str(exponent)
-    if exponent < 0: return sign + digits + "e" + str(exponent)
+    return sign + digits + sep + "{:+}".format(exponent)
 
 def str_to_man_exp(x, base=10):
     """Helper function for from_str."""

mpmath/tests/test_convert.py

@@ -61,6 +61,13 @@ def test_from_str():
     assert mpf(from_str('0b1101.100101')) == mpf('13.578125')
     assert mpf(from_str('0o1101.100101')) == mpf('577.12524795532227')
 
+def test_to_str():
+    assert to_str(from_str('ABC.ABC', base=16), 6, base=16) == '0xabc.abc'
+    assert to_str(from_str('0x3.a7p10', base=16), 3, base=16) == '0xe9c.0'
+    assert to_str(from_str('0x1.4ace478p+33'), 7, base=16) == '0x2.959c8f@+8'
+    assert to_str(from_str('0o1101.100101'), 8, base=8) == '0o1101.1001'
+    assert to_str(from_str('0b1101.100101'), 10, base=2) == '0b1101.100101'
+
 def test_pretty():
     mp.pretty = True
     assert repr(mpf(2.5)) == '2.5'