binary_repr handles negative numbers now, and takes an optional width argument.

base_repr raise ValueError if the number is negative or base > 36

2 files changed, 31 insertions, 13 deletions

diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index 6ca47a87b..f98ca9af0 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -546,29 +546,44 @@ _lkup = {
     'F':'1111',
     'L':''}
 
-def binary_repr(num):
+def binary_repr(num, width=None):
     """Return the binary representation of the input number as a string.
 
     This is equivalent to using base_repr with base 2, but about 25x
     faster.
+
+    For negative numbers, if width is not given, a - sign is added to the
+    front. If width is given, the two's complement of the number is
+    returned, with respect to that width.
     """
-    ostr = hex(num)
-    bin = ''
-    for ch in ostr[2:]:
-        bin += _lkup[ch]
-    if '1' in bin:
-        ind = 0
-        while bin[ind] == '0':
-            ind += 1
-        return bin[ind:]
-    else:
+    sign = ''
+    if num < 0:
+        if width is None:
+            sign = '-'
+            num = -num
+        else:
+            # replace num with its 2-complement
+            num = 2**width + num
+    elif num == 0:
         return '0'
+    ostr = hex(num)
+    bin = ''.join([_lkup[ch] for ch in ostr[2:]])
+    bin = bin.lstrip('0')
+    if width is not None:
+        bin = bin.zfill(width)
+    return sign + bin
 
 def base_repr (number, base=2, padding=0):
-    """Return the representation of a number in any given base.
+    """Return the representation of a number in the given base.
+
+    Base can't be larger than 36.
     """
-    chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    if number < 0:
+        raise ValueError("negative numbers not handled in base_repr")
+    if base > 36:
+        raise ValueError("bases greater than 36 not handled in base_repr")
 
+    chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
     import math
     lnb = math.log(base)
     res = padding*chars[0]
diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py
index acda7b081..f4c4431b6 100644
--- a/numpy/core/tests/test_numeric.py
+++ b/numpy/core/tests/test_numeric.py
@@ -248,6 +248,9 @@ class test_binary_repr(NumpyTestCase):
     def test_large(self):
         assert_equal(binary_repr(10736848),'101000111101010011010000')
 
+    def test_negative(self):
+        assert_equal(binary_repr(-1), '-1')
+        assert_equal(binary_repr(-1, width=8), '11111111')
 
 def assert_array_strict_equal(x, y):
     assert_array_equal(x, y)