1 files changed, 89 insertions, 2 deletions
diff --git a/numpy/polynomial/_polybase.py b/numpy/polynomial/_polybase.py
index 78392d2a2..9f4d30e53 100644
--- a/numpy/polynomial/_polybase.py
+++ b/numpy/polynomial/_polybase.py
@@ -9,7 +9,7 @@ abc module from the stdlib, hence it is only available for Python >= 2.6.
 from __future__ import division, absolute_import, print_function
 
 from abc import ABCMeta, abstractmethod, abstractproperty
-from numbers import Number
+import numbers
 
 import numpy as np
 from . import polyutils as pu
@@ -82,6 +82,10 @@ class ABCPolyBase(object):
     def nickname(self):
         pass
 
+    @abstractproperty
+    def basis_name(self):
+        pass
+
     @abstractmethod
     def _add(self):
         pass
@@ -273,6 +277,89 @@ class ABCPolyBase(object):
         name = self.nickname
         return format % (name, coef)
 
+    @classmethod
+    def _repr_latex_term(cls, i, arg_str, needs_parens):
+        if cls.basis_name is None:
+            raise NotImplementedError(
+                "Subclasses must define either a basis name, or override "
+                "_repr_latex_term(i, arg_str, needs_parens)")
+        # since we always add parens, we don't care if the expression needs them
+        return "{{{basis}}}_{{{i}}}({arg_str})".format(
+            basis=cls.basis_name, i=i, arg_str=arg_str
+        )
+
+    @staticmethod
+    def _repr_latex_scalar(x):
+        # TODO: we're stuck with disabling math formatting until we handle
+        # exponents in this function
+        return r'\text{{{}}}'.format(x)
+
+    def _repr_latex_(self):
+        # get the scaled argument string to the basis functions
+        off, scale = self.mapparms()
+        if off == 0 and scale == 1:
+            term = 'x'
+            needs_parens = False
+        elif scale == 1:
+            term = '{} + x'.format(
+                self._repr_latex_scalar(off)
+            )
+            needs_parens = True
+        elif off == 0:
+            term = '{}x'.format(
+                self._repr_latex_scalar(scale)
+            )
+            needs_parens = True
+        else:
+            term = '{} + {}x'.format(
+                self._repr_latex_scalar(off),
+                self._repr_latex_scalar(scale)
+            )
+            needs_parens = True
+
+        # filter out uninteresting coefficients
+        filtered_coeffs = [
+            (i, c)
+            for i, c in enumerate(self.coef)
+            # if not (c == 0)  # handle NaN
+        ]
+
+        mute = r"\color{{LightGray}}{{{}}}".format
+
+        parts = []
+        for i, c in enumerate(self.coef):
+            # prevent duplication of + and - signs
+            if i == 0:
+                coef_str = '{}'.format(self._repr_latex_scalar(c))
+            elif not isinstance(c, numbers.Real):
+                coef_str = ' + ({})'.format(self._repr_latex_scalar(c))
+            elif not np.signbit(c):
+                coef_str = ' + {}'.format(self._repr_latex_scalar(c))
+            else:
+                coef_str = ' - {}'.format(self._repr_latex_scalar(-c))
+
+            # produce the string for the term
+            term_str = self._repr_latex_term(i, term, needs_parens)
+            if term_str == '1':
+                part = coef_str
+            else:
+                part = r'{}\,{}'.format(coef_str, term_str)
+
+            if c == 0:
+                part = mute(part)
+
+            parts.append(part)
+
+        if parts:
+            body = ''.join(parts)
+        else:
+            # in case somehow there are no coefficients at all
+            body = '0'
+
+        return r'$x \mapsto {}$'.format(body)
+
+
+
     # Pickle and copy
 
     def __getstate__(self):
@@ -338,7 +425,7 @@ class ABCPolyBase(object):
         # there is no true divide if the rhs is not a Number, although it
         # could return the first n elements of an infinite series.
         # It is hard to see where n would come from, though.
-        if not isinstance(other, Number) or isinstance(other, bool):
+        if not isinstance(other, numbers.Number) or isinstance(other, bool):
             form = "unsupported types for true division: '%s', '%s'"
             raise TypeError(form % (type(self), type(other)))
         return self.__floordiv__(other)