catacomb/__init__.py: Support IPython's pretty-printer.

Add `_repr_pretty_' methods to many types to improve presentation.
Also, sneakily add a bunch of printing methods to the key-management
classes.

diff --git a/catacomb/__init__.py b/catacomb/__init__.py
index 022c5f97119fc6e88b73a834e529a5d3dab5a14b..d0478ce78bcc4bf3bf4c2e9005f9745187bcba44 100644 (file)
--- a/catacomb/__init__.py
+++ b/catacomb/__init__.py
@@ -82,6 +82,31 @@ def _checkend(r):
     raise SyntaxError, 'junk at end of string'
   return x
 
+## Some pretty-printing utilities.
+def _pp_str(me, pp, cyclep): pp.text(cyclep and '...' or str(me))
+def _pp_kv(pp, k, v):
+  ind = len(k) + 3
+  pp.begin_group(ind, k + ' = ')
+  pp.pretty(v)
+  pp.end_group(ind, '')
+def _pp_commas(pp, printfn, items):
+  firstp = True
+  for i in items:
+    if firstp: firstp = False
+    else: pp.text(','); pp.breakable()
+    printfn(i)
+def _pp_dict(pp, items):
+  def p((k, v)):
+    pp.begin_group(0, '')
+    pp.pretty(k)
+    pp.text(':')
+    pp.begin_group(2, '')
+    pp.breakable()
+    pp.pretty(v)
+    pp.end_group(2, '')
+    pp.end_group(0, '')
+  _pp_commas(pp, p, items)
+
 ###--------------------------------------------------------------------------
 ### Bytestrings.
 
@@ -150,6 +175,7 @@ class BaseRat (object):
   def denom(me): return me._d
   def __str__(me): return '%s/%s' % (me._n, me._d)
   def __repr__(me): return '%s(%s, %s)' % (type(me).__name__, me._n, me._d)
+  _repr_pretty_ = _pp_str
 
   def __add__(me, you):
     n, d = _split_rat(you)
@@ -194,6 +220,7 @@ class _tmp:
   def reduce(x): return MPReduce(x)
   def __div__(me, you): return IntRat(me, you)
   def __rdiv__(me, you): return IntRat(you, me)
+  _repr_pretty_ = _pp_str
 _augment(MP, _tmp)
 
 class _tmp:
@@ -205,6 +232,7 @@ class _tmp:
   def quadsolve(x, y): return x.reduce().quadsolve(y)
   def __div__(me, you): return GFRat(me, you)
   def __rdiv__(me, you): return GFRat(you, me)
+  _repr_pretty_ = _pp_str
 _augment(GF, _tmp)
 
 class _tmp:
@@ -225,12 +253,30 @@ _augment(Field, _tmp)
 class _tmp:
   def __repr__(me): return '%s(%sL)' % (type(me).__name__, me.p)
   def __hash__(me): return 0x114401de ^ hash(me.p)
+  def _repr_pretty_(me, pp, cyclep):
+    tyname = type(me).__name__
+    if cyclep:
+      pp.text(tyname + '(...)')
+    else:
+      ind = len(tyname) + 1
+      pp.begin_group(ind, tyname + '(')
+      pp.pretty(me.p)
+      pp.end_group(ind, ')')
   def ec(me, a, b): return ECPrimeProjCurve(me, a, b)
 _augment(PrimeField, _tmp)
 
 class _tmp:
   def __repr__(me): return '%s(%#xL)' % (type(me).__name__, me.p)
   def ec(me, a, b): return ECBinProjCurve(me, a, b)
+  def _repr_pretty_(me, pp, cyclep):
+    tyname = type(me).__name__
+    if cyclep:
+      pp.text(tyname + '(...)')
+    else:
+      ind = len(tyname) + 1
+      pp.begin_group(ind, tyname + '(')
+      pp.text('%#x' % me.p)
+      pp.end_group(ind, ')')
 _augment(BinField, _tmp)
 
 class _tmp:
@@ -248,6 +294,7 @@ _augment(BinNormField, _tmp)
 class _tmp:
   def __str__(me): return str(me.value)
   def __repr__(me): return '%s(%s)' % (repr(me.field), repr(me.value))
+  _repr_pretty_ = _pp_str
 _augment(FE, _tmp)
 
 ###--------------------------------------------------------------------------
@@ -256,6 +303,17 @@ _augment(FE, _tmp)
 class _tmp:
   def __repr__(me):
     return '%s(%r, %s, %s)' % (type(me).__name__, me.field, me.a, me.b)
+  def _repr_pretty_(me, pp, cyclep):
+    tyname = type(me).__name__
+    if cyclep:
+      pp.text(tyname + '(...)')
+    else:
+      ind = len(tyname) + 1
+      pp.begin_group(ind, tyname + '(')
+      pp.pretty(me.field); pp.text(','); pp.breakable()
+      pp.pretty(me.a); pp.text(','); pp.breakable()
+      pp.pretty(me.b)
+      pp.end_group(ind, ')')
   def frombuf(me, s):
     return ecpt.frombuf(me, s)
   def fromraw(me, s):
@@ -289,12 +347,32 @@ class _tmp:
   def __str__(me):
     if not me: return 'inf'
     return '(%s, %s)' % (me.ix, me.iy)
+  def _repr_pretty_(me, pp, cyclep):
+    if cyclep:
+      pp.text('...')
+    elif not me:
+      pp.text('inf')
+    else:
+      pp.begin_group(1, '(')
+      pp.pretty(me.ix); pp.text(','); pp.breakable()
+      pp.pretty(me.iy)
+      pp.end_group(1, ')')
 _augment(ECPt, _tmp)
 
 class _tmp:
   def __repr__(me):
     return 'ECInfo(curve = %r, G = %r, r = %s, h = %s)' % \
            (me.curve, me.G, me.r, me.h)
+  def _repr_pretty_(me, pp, cyclep):
+    if cyclep:
+      pp.text('ECInfo(...)')
+    else:
+      pp.begin_group(7, 'ECInfo(')
+      _pp_kv(pp, 'curve', me.curve); pp.text(','); pp.breakable()
+      _pp_kv(pp, 'G', me.G); pp.text(','); pp.breakable()
+      _pp_kv(pp, 'r', me.r); pp.text(','); pp.breakable()
+      _pp_kv(pp, 'h', me.h)
+      pp.end_group(7, ')')
   def __hash__(me):
     h = 0x9bedb8de
     h ^=   hash(me.curve)
@@ -326,6 +404,16 @@ class _tmp:
   def __repr__(me):
     return 'KeySZRange(%d, %d, %d, %d)' % \
            (me.default, me.min, me.max, me.mod)
+  def _repr_pretty_(me, pp, cyclep):
+    if cyclep:
+      pp.text('KeySZRange(...)')
+    else:
+      pp.begin_group(11, 'KeySZRange(')
+      pp.pretty(me.default); pp.text(','); pp.breakable()
+      pp.pretty(me.min); pp.text(','); pp.breakable()
+      pp.pretty(me.max); pp.text(','); pp.breakable()
+      pp.pretty(me.mod)
+      pp.end_group(11, ')')
   def check(me, sz): return me.min <= sz <= me.max and sz % me.mod == 0
   def best(me, sz):
     if sz < me.min: raise ValueError, 'key too small'
@@ -335,6 +423,16 @@ _augment(KeySZRange, _tmp)
 
 class _tmp:
   def __repr__(me): return 'KeySZSet(%d, %s)' % (me.default, me.set)
+  def _repr_pretty_(me, pp, cyclep):
+    if cyclep:
+      pp.text('KeySZSet(...)')
+    else:
+      pp.begin_group(9, 'KeySZSet(')
+      pp.pretty(me.default); pp.text(','); pp.breakable()
+      pp.begin_group(1, '{')
+      _pp_commas(pp, pp.pretty, me.set)
+      pp.end_group(1, '}')
+      pp.end_group(9, ')')
   def check(me, sz): return sz in me.set
   def best(me, sz):
     found = -1
@@ -344,6 +442,104 @@ class _tmp:
     return found
 _augment(KeySZSet, _tmp)
 
+###--------------------------------------------------------------------------
+### Key data objects.
+
+class _tmp:
+  def __repr__(me): return 'KeyFile(%r)' % me.name
+_augment(KeyFile, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'Key(%r)' % me.fulltag
+_augment(Key, _tmp)
+
+class _tmp:
+  def __repr__(me):
+    return 'KeyAttributes({%s})' % \
+        ', '.join(['%r: %r' % kv for kv in me.iteritems()])
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(14, 'KeyAttributes(')
+    if cyclep: pp.text('...')
+    else: _pp_dict(pp, me.iteritems())
+    pp.end_group(14, ')')
+_augment(KeyAttributes, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'KeyDataBinary(%r, %r)' % \
+        (me.bin, me.writeflags(me.flags))
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(14, 'KeyDataBinary(')
+    if cyclep:
+      pp.text('KeyDataBinary(...)')
+    else:
+      pp.pretty(me.bin); pp.text(','); pp.breakable()
+      pp.pretty(me.writeflags(me.flags))
+    pp.end_group(14, ')')
+_augment(KeyDataBinary, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'KeyDataEncrypted(%r, %r)' % \
+        (me.ct, me.writeflags(me.flags))
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(17, 'KeyDataEncrypted(')
+    if cyclep:
+      pp.text('...')
+    else:
+      pp.pretty(me.ct); pp.text(','); pp.breakable()
+      pp.pretty(me.writeflags(me.flags))
+    pp.end_group(17, ')')
+_augment(KeyDataEncrypted, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'KeyDataMP(%r, %r)' % \
+        (me.mp, me.writeflags(me.flags))
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(10, 'KeyDataMP(')
+    if cyclep:
+      pp.text('...')
+    else:
+      pp.pretty(me.mp); pp.text(','); pp.breakable()
+      pp.pretty(me.writeflags(me.flags))
+    pp.end_group(10, ')')
+_augment(KeyDataMP, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'KeyDataString(%r)' % \
+        (me.str, me.writeflags(me.flags))
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(14, 'KeyDataString(')
+    if cyclep:
+      pp.text('...')
+    else:
+      pp.pretty(me.str); pp.text(','); pp.breakable()
+      pp.pretty(me.writeflags(me.flags))
+    pp.end_group(14, ')')
+_augment(KeyDataString, _tmp)
+
+class _tmp:
+  def __repr__(me): return 'KeyDataECPt(%r)' % \
+        (me.ecpt, me.writeflags(me.flags))
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(12, 'KeyDataECPt(')
+    if cyclep:
+      pp.text('...')
+    else:
+      pp.pretty(me.ecpt); pp.text(','); pp.breakable()
+      pp.pretty(me.writeflags(me.flags))
+    pp.end_group(12, ')')
+_augment(KeyDataECPt, _tmp)
+
+class _tmp:
+  def __repr__(me):
+    return 'KeyDataStructured({%s})' % \
+        ', '.join(['%r: %r' % kv for kv in me.iteritems()])
+  def _repr_pretty_(me, pp, cyclep):
+    pp.begin_group(20, 'KeyDataStructured({ ')
+    if cyclep: pp.text('...')
+    else: _pp_dict(pp, me.iteritems())
+    pp.end_group(20, ' })')
+_augment(KeyDataStructured, _tmp)
+
 ###--------------------------------------------------------------------------
 ### Abstract groups.
 
@@ -351,6 +547,17 @@ class _tmp:
   def __repr__(me):
     return '%s(p = %s, r = %s, g = %s)' % \
            (type(me).__name__, me.p, me.r, me.g)
+  def _repr_pretty_(me, pp, cyclep):
+    tyname = type(me).__name__
+    if cyclep:
+      pp.text(tyname + '(...)')
+    else:
+      ind = len(tyname) + 1
+      pp.begin_group(ind, tyname + '(')
+      _pp_kv(pp, 'p', me.p); pp.text(','); pp.breakable()
+      _pp_kv(pp, 'r', me.r); pp.text(','); pp.breakable()
+      _pp_kv(pp, 'g', me.g)
+      pp.end_group(ind, ')')
 _augment(FGInfo, _tmp)
 
 class _tmp:
@@ -393,6 +600,7 @@ _augment(ECGroup, _tmp)
 class _tmp:
   def __repr__(me):
     return '%r(%r)' % (me.group, str(me))
+  _repr_pretty_ = _pp_str
 _augment(GE, _tmp)
 
 ###--------------------------------------------------------------------------
@@ -534,7 +742,14 @@ class _groupmap (object):
     me._n = max(map.values()) + 1
     me.i = me._n*[None]
   def __repr__(me):
-    return '{%s}' % ', '.join(['%r: %r' % (k, me[k]) for k in me])
+    return '{%s}' % ', '.join(['%r: %r' % kv for kv in me.iteritems()])
+  def _repr_pretty_(me, pp, cyclep):
+    if cyclep:
+      pp.text('{ ... }')
+    else:
+      pp.begin_group(2, '{ ')
+      _pp_dict(pp, me.iteritems())
+      pp.end_group(2, ' }')
   def __len__(me):
     return me._n
   def __contains__(me, k):