sympy/printing/pretty.py

   1 from sympy.core import Basic, Function  # XXX remove Function from here
   2 from printer import Printer
   3 from stringpict import *
   4
   5 import re
   6
   7 _use_unicode = False
   8
   9 def pprint_use_unicode(flag = None):
  10     """Set whether pretty-printer should use unicode by default"""
  11     global _use_unicode
  12     if flag is None:
  13         return _use_unicode
  14
  15     use_unicode_prev = _use_unicode
  16     _use_unicode = flag
  17     return use_unicode_prev
  18
  19 class PrettyPrinter(Printer):
  20     """
  21     A class that prints a prettified expression, one that is not limited
  22     to one dimension like casting the expression to a string would return.
  23     """
  24     def __init__(self, use_unicode=None):
  25         Printer.__init__(self)
  26         if use_unicode is None:
  27             use_unicode = _use_unicode
  28         self._use_unicode = use_unicode
  29         if self._use_unicode:
  30             self._str = unicode
  31         self.emptyPrinter = lambda x : prettyForm(self._str(x), prettyForm.ATOM)
  32
  33     def doprint(self, expr):
  34         Printer.doprint.__doc__
  35         return self._print(expr).terminal_string()
  36
  37     def _print_Symbol(self, e):
  38         if self._use_unicode:
  39             # letter: (uppercase, lowercase)
  40             greek = {
  41                 'alpha': (u'\u0391', u'\u03b1'),
  42                 'beta': (u'\u0392', u'\u03b2'),
  43                 'gamma': (u'\u0393', u'\u03b3'),
  44                 'delta': (u'\u0394', u'\u03b4'),
  45                 'epsilon': (u'\u0395', u'\u03b5'),
  46                 'zeta': (u'\u0396', u'\u03b6'),
  47                 'eta': (u'\u0397', u'\u03b7'),
  48                 'theta': (u'\u0398', u'\u03b8'),
  49                 'iota': (u'\u0399', u'\u03b9'),
  50                 'kappa': (u'\u039a', u'\u03ba'),
  51                 'lambda': (u'\u039b', u'\u03bb'),
  52                 'mu': (u'\u039c', u'\u03bc'),
  53                 'nu': (u'\u039d', u'\u03bd'),
  54                 'xi': (u'\u039e', u'\u03be'),
  55                 'omicron': (u'\u039f', u'\u03bf'),
  56                 'pi': (u'\u03a0', u'\u03c0'),
  57                 'rho': (u'\u03a1', u'\u03c1'),
  58                 'sigma': (u'\u03a3', u'\u03c3'),
  59                 'tau': (u'\u03a4', u'\u03c4'),
  60                 'upsilon': (u'\u03a5', u'\u03c5'),
  61                 'phi': (u'\u03a6', u'\u03c6'),
  62                 'chi': (u'\u03a7', u'\u03c7'),
  63                 'psi': (u'\u03a8', u'\u03c8'),
  64                 'omega': (u'\u03a9', u'\u03c9')
  65             }
  66
  67             name = e.name.lower()
  68
  69             # let's split name into symbol + index
  70             # UC: beta1
  71             m = re.match('(\D+)(\d*)$', name)
  72             if m is None:
  73                 return
  74
  75             name = m.expand(r'\1')
  76             idx  = m.expand(r'\2')
  77
  78             if name in greek:
  79                 # If first character lowercase, use lowercase greek letter
  80                 if name[0] == e.name[0]:
  81                     greek_name = greek[name][1]
  82                 else:
  83                     greek_name = greek[name][0]
  84
  85                 return prettyForm(greek_name+idx, binding=prettyForm.ATOM)
  86
  87     def _print_Factorial(self, e):
  88         x = e[0]
  89         if (isinstance(x, Basic.Integer) and x.is_nonnegative) or \
  90             isinstance(x, Basic.Symbol):
  91             s = self._print(x)
  92         else:
  93             s = "(" + self._print(x) + ")"
  94         return s + "!"
  95
  96     def _print_Exp1(self, e):
  97         if self._use_unicode:
  98             return prettyForm(u'\u212f', binding=prettyForm.ATOM)
  99
 100     def _print_Pi(self, e):
 101         if self._use_unicode:
 102             return prettyForm(u'\u03c0', binding=prettyForm.ATOM)
 103
 104     def _print_Infinity(self, e):
 105         if self._use_unicode:
 106             return prettyForm(u'\u221e', binding=prettyForm.ATOM)
 107
 108     def _print_NegativeInfinity(self, e):
 109         if self._use_unicode:
 110             return prettyForm(u'-\u221e', binding=prettyForm.ATOM)
 111
 112     def _print_ImaginaryUnit(self, e):
 113         if self._use_unicode:
 114             return prettyForm(u'\u03b9', binding=prettyForm.ATOM)
 115
 116     def _print_Relational(self, e):
 117         charmap = {
 118             '==': ('=', '='),
 119             '<':  ('<', '<'),
 120             '<=': ('<=', u'\u2264'),
 121             '!=': ('!=', u'\u2260')
 122         }
 123         if self._use_unicode:
 124             op = charmap[e.rel_op][1]
 125         else:
 126             op = charmap[e.rel_op][0]
 127         op = prettyForm(' ' + op + ' ')
 128
 129         l = self._print(e.lhs)
 130         r = self._print(e.rhs)
 131         pform = prettyForm(*stringPict.next(l, op))
 132         pform = prettyForm(*stringPict.next(pform, r))
 133         return pform
 134
 135     def _print_ApplyConjugate(self, e):
 136         pform = self._print(e[0])
 137         return prettyForm(*stringPict.above(pform, '_'*pform.width()))
 138
 139     def _print_abs(self, e):
 140         pform = self._print(e[0])
 141         pform.baseline = 0
 142         bars = '|' + ('\n|' * (pform.height()-1))
 143         pform = prettyForm(*stringPict.next(bars, pform))
 144         pform = prettyForm(*stringPict.next(pform, bars))
 145         return pform
 146
 147     def _print_Derivative(self, deriv):
 148         syms = list(deriv.symbols)
 149         syms.reverse()
 150         x = None
 151         for sym in syms:
 152             if x is None:
 153                 x = prettyForm('d' + str(sym))
 154             else:
 155                 x = prettyForm(*stringPict.next(x, ' d' + str(sym)))
 156
 157         f = prettyForm(binding=prettyForm.FUNC, *self._print(deriv.expr).parens())
 158
 159         pform = prettyForm('d')
 160         if len(syms) > 1:
 161             pform = pform ** prettyForm(str(len(deriv.symbols)))
 162
 163         pform = prettyForm(*pform.below(stringPict.LINE, x))
 164         pform.baseline = pform.baseline + 1
 165         pform = prettyForm(*stringPict.next(pform, f))
 166         return pform
 167
 168     def _print_Integral(self, integral):
 169         f   = integral.function
 170
 171         # Add parentheses if a sum and create pretty form for argument
 172         prettyF = self._print(f)
 173         if isinstance(f, Basic.Add):
 174             prettyF = prettyForm(*prettyF.parens())
 175
 176         # dx dy dz ...
 177         arg = prettyF
 178         for x,ab in integral.limits:
 179             prettyArg = self._print(x)
 180             if prettyArg.width() > 1:
 181                 arg = prettyForm(*arg.right(' d(', prettyArg, ')'))
 182             else:
 183                 arg = prettyForm(*arg.right(' d', prettyArg))
 184
 185         arg.baseline = 0
 186
 187         # \int \int \int ...
 188         firstterm = True
 189         S = None
 190         for x,ab in integral.limits:
 191             # Create bar based on the height of the argument
 192             bar = ' |' + ('\r |' * (arg.height()+1))
 193
 194             # Construct the pretty form with the integral sign and the argument
 195             pform = prettyForm(bar)
 196             pform = prettyForm(*pform.below(' /  '))
 197             pform = prettyForm(*pform.above('   /'))
 198             pform.baseline = (arg.height() + 3)/2
 199
 200             if ab is not None:
 201                 # Create pretty forms for endpoints, if definite integral
 202                 prettyA = self._print(ab[0])
 203                 prettyB = self._print(ab[1])
 204
 205                 # Add spacing so that endpoint can more easily be
 206                 # identified with the correct integral sign
 207                 spc = max(1, 4 - prettyB.width())
 208                 prettyB = prettyForm(*prettyB.left(' ' * spc))
 209
 210                 pform = prettyForm(*pform.above(prettyB))
 211                 pform = prettyForm(*pform.below(prettyA))
 212                 pform = prettyForm(*pform.right(' '))
 213
 214             if firstterm:
 215                 S = pform   # first term
 216                 firstterm = False
 217             else:
 218                 S = prettyForm(*S.left(pform))
 219
 220         pform = prettyForm(*arg.left(S))
 221         return pform
 222
 223     def _print_exp(self, e):
 224         if self._use_unicode:
 225             base = prettyForm(u'\u212f', binding=prettyForm.ATOM)
 226         else:
 227             base = prettyForm('e', binding=prettyForm.ATOM)
 228         return base ** self._print(e[0])
 229
 230     # TODO: rename to _print_Function
 231     #       after 'Function -> Function' transition is done
 232     def _print_Function(self, e):
 233         return self._print_Apply(e) # XXX this is temporary
 234
 235     def _print_Apply(self, e):
 236         func = e.func
 237         args = e[:]
 238         n = len(args)
 239
 240         # XXX temp. hack -- remove when 'Function -> Function' is done
 241         #if isinstance(func, Function):
 242         #    func_name = func.name
 243         #else:
 244         func_name = func.__name__
 245
 246         prettyFunc = self._print(Basic.Symbol(func_name));
 247         prettyArgs = self._print(args[0])
 248         for i in xrange(1, n):
 249             pform = self._print(args[i])
 250             prettyArgs = prettyForm(*stringPict.next(prettyArgs, ', '))
 251             prettyArgs = prettyForm(*stringPict.next(prettyArgs, pform))
 252
 253         prettyArgs = prettyForm(*prettyArgs.left('('))
 254         prettyArgs = prettyForm(*prettyArgs.right(')'))
 255
 256         pform = prettyForm(binding=prettyForm.FUNC, *stringPict.next(prettyFunc, prettyArgs))
 257
 258         # store pform parts so it can be reassembled e.g. when powered
 259         pform.prettyFunc = prettyFunc
 260         pform.prettyArgs = prettyArgs
 261
 262         return pform
 263
 264     def _print_Add(self, sum):
 265         pforms = []
 266         for x in sum:
 267             # Check for negative "things" so that this information can be enforce upon
 268             # the pretty form so that it can be made of use (such as in a sum).
 269             if isinstance(x, Basic.Mul) and x.as_coeff_terms()[0] < 0:
 270                 pform1 = self._print(-x)
 271                 if len(pforms) == 0:
 272                     if pform1.height() > 1:
 273                         pform2 = '- '
 274                     else:
 275                         pform2 = '-'
 276                 else:
 277                     pform2 = ' - '
 278                 pform = stringPict.next(pform2, pform1)
 279                 pforms.append(prettyForm(binding=prettyForm.NEG, *pform))
 280             elif isinstance(x, Basic.Number) and x < 0:
 281                 pform1 = self._print(-x)
 282                 if len(pforms) == 0:
 283                     if pform1.height() > 1:
 284                         pform2 = '- '
 285                     else:
 286                         pform2 = '-'
 287                     pform = stringPict.next(pform2, pform1)
 288                 else:
 289                     pform = stringPict.next(' - ', pform1)
 290                 pforms.append(prettyForm(binding=prettyForm.NEG, *pform))
 291             else:
 292                 pforms.append(self._print(x))
 293         return prettyForm.__add__(*pforms)
 294
 295     def _print_Mul(self, product):
 296         a = [] # items in the numerator
 297         b = [] # items that are in the denominator (if any)
 298
 299         # Gather terms for numerator/denominator
 300         for item in product:
 301             if isinstance(item, Basic.Pow) and item.exp == -1:
 302                 b.append(item.base)
 303             elif isinstance(item, Basic.Rational):
 304                 if item.p != 1:
 305                     a.append( Basic.Rational(item.p) )
 306                 if item.q != 1:
 307                     b.append( Basic.Rational(item.q) )
 308             else:
 309                 a.append(item)
 310
 311         # Convert to pretty forms. Add parens to Add instances if there
 312         # is more than one term in the numer/denom
 313         for i in xrange(0, len(a)):
 314             if isinstance(a[i], Basic.Add) and len(a) > 1:
 315                 a[i] = prettyForm(*self._print(a[i]).parens())
 316             else:
 317                 a[i] = self._print(a[i])
 318
 319         for i in xrange(0, len(b)):
 320             if isinstance(b[i], Basic.Add) and len(b) > 1:
 321                 b[i] = prettyForm(*self._print(b[i]).parens())
 322             else:
 323                 b[i] = self._print(b[i])
 324
 325         # Construct a pretty form
 326         if len(b) == 0:
 327             return prettyForm.__mul__(*a)
 328         else:
 329             if len(a) == 0:
 330                 a.append( self._print(Basic.One()) )
 331             return prettyForm.__mul__(*a) / prettyForm.__mul__(*b)
 332
 333     def _print_Pow(self, power):
 334         if isinstance(power.exp, Basic.Half):
 335             # If it's a square root
 336             bpretty = self._print(power.base)
 337             bl = int((bpretty.height() / 2.0) + 0.5)
 338
 339             s2 = stringPict("\\/")
 340             for x in xrange(1, bpretty.height()):
 341                 s3 = stringPict(" " * (2*x+1) + "/")
 342                 s2 = stringPict(*s2.above(s3))
 343             s2.baseline = -1
 344
 345             s = prettyForm("__" + "_" * bpretty.width())
 346             s = prettyForm(*stringPict.below(s, bpretty))
 347             s = prettyForm(*stringPict.left(s, s2))
 348             s.baseline = bl
 349             return s
 350         elif power.exp == -1:
 351             # Things like 1/x
 352             return prettyForm("1") / self._print(power.base)
 353
 354         # None of the above special forms, do a standard power
 355         b,e = power.as_base_exp()
 356         return self._print(b)**self._print(e)
 357
 358     def _print_Rational(self, r):
 359         if r.q == 1:
 360             return prettyForm(str(r.p), prettyForm.ATOM)
 361         elif abs(r.p) > 10 and abs(r.q) > 10:
 362             # If more than one digit in numer and denom, print larger fraction
 363             if r.is_negative:
 364                 pform = prettyForm(str(-r.p))/prettyForm(str(r.q))
 365                 return prettyForm(binding=prettyForm.NEG, *pform.left('- '))
 366             else:
 367                 return prettyForm(str(r.p))/prettyForm(str(r.q))
 368
 369 def pretty(expr, use_unicode=None):
 370     """
 371     Returns a string containing the prettified form of expr. If use_unicode
 372     is set to True then certain expressions will use unicode characters,
 373     such as the greek letter pi for Basic.Pi instances.
 374     """
 375     pp = PrettyPrinter(use_unicode)
 376     return pp.doprint(expr)
 377
 378 def pretty_print(expr, use_unicode=None):
 379     """
 380     Prints expr in pretty form.
 381
 382     pprint is just a shortcut for this function
 383     """
 384     print pretty(expr, use_unicode)
 385
 386 pprint = pretty_print