Repository: numba/numba

Issue:
Support signature kwarg in dispatcher `.inspect_types` method
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Feature request

<!--

Please include details of the feature you would like to see, why you would
like to see it/the use case

-->

Just noticed that the dispatcher object `.inspect_types` method doesn't support signatures cf. all the other `.inspect_` methods, this ought to be implemented. e.g.
```pycon
In [13]: from numba import njit                                                                                                 

In [14]: @njit 
    ...: def foo(x): 
    ...:     return x + 1 
    ...:                                                                                                                        

In [15]: foo(10)                                                                                                                
Out[15]: 11

In [16]: foo.inspect_types(foo.signatures[0])  # <--- doesn't work
```



Discussion:

Resolving patch:
diff --git a/numba/dispatcher.py b/numba/dispatcher.py
--- a/numba/dispatcher.py
+++ b/numba/dispatcher.py
@@ -383,7 +383,7 @@ def inspect_asm(self, signature=None):
 
         return dict((sig, self.inspect_asm(sig)) for sig in self.signatures)
 
-    def inspect_types(self, file=None, **kwargs):
+    def inspect_types(self, file=None, signature=None, **kwargs):
         """
         print or return annotated source with Numba intermediate IR
 
@@ -394,11 +394,15 @@ def inspect_types(self, file=None, **kwargs):
         pretty = kwargs.get('pretty', False)
         style = kwargs.get('style', 'default')
 
+        overloads = self.overloads
+        if signature is not None:
+            overloads = {signature: self.overloads[signature]}
+
         if not pretty:
             if file is None:
                 file = sys.stdout
 
-            for ver, res in utils.iteritems(self.overloads):
+            for ver, res in utils.iteritems(overloads):
                 print("%s %s" % (self.py_func.__name__, ver), file=file)
                 print('-' * 80, file=file)
                 print(res.type_annotation, file=file)
@@ -407,7 +411,7 @@ def inspect_types(self, file=None, **kwargs):
             if file is not None:
                 raise ValueError("`file` must be None if `pretty=True`")
             from .pretty_annotate import Annotate
-            return Annotate(self, style=style)
+            return Annotate(self, signature=signature, style=style)
 
     def inspect_cfg(self, signature=None, show_wrapper=None):
         """
@@ -437,10 +441,15 @@ def get_annotation_info(self, signature=None):
         signature. If no signature is supplied a dictionary of signature to
         annotation information is returned.
         """
-        if signature is not None:
-            cres = self.overloads[signature]
-            return cres.type_annotation.annotate_raw()
-        return dict((sig, self.annotate(sig)) for sig in self.signatures)
+        signatures = self.signatures if signature is None else [signature]
+        out = collections.OrderedDict()
+        for sig in signatures:
+            cres = self.overloads[sig]
+            ta = cres.type_annotation
+            key = (ta.func_id.filename + ':' + str(ta.func_id.firstlineno + 1),
+                   ta.signature)
+            out[key] = ta.annotate_raw()[key]
+        return out
 
     def _explain_ambiguous(self, *args, **kws):
         """
diff --git a/numba/pretty_annotate.py b/numba/pretty_annotate.py
--- a/numba/pretty_annotate.py
+++ b/numba/pretty_annotate.py
@@ -228,31 +228,46 @@ class Annotate:
     Example:
 
     >>> import numba
-    ... from numba.pretty_annotate import Annotate
-    ...
-    ... @numba.jit
+    >>> from numba.pretty_annotate import Annotate
+    >>> @numba.jit
     ... def test(q):
     ...     res = 0
     ...     for i in range(q):
     ...         res += i
     ...     return res
     ...
-    ... test(10)
-    ... Annotate(test)
+    >>> test(10)
+    45
+    >>> Annotate(test)
 
     The last line will return an HTML and/or ANSI representation that will be
     displayed accordingly in Jupyter/IPython.
 
     Function annotations persist across compilation for newly encountered
-    type signatures and as a result annotations are shown for all signatures.
+    type signatures and as a result annotations are shown for all signatures
+    by default.
+
+    Annotations for a specific signature can be shown by using the
+    ``signature`` parameter.
+
+    >>> @numba.jit
+    ... def add(x, y):
+    ...     return x + y
+    ...
+    >>> add(1, 2)
+    3
+    >>> add(1.3, 5.7)
+    7.0
+    >>> add.signatures
+    [(int64, int64), (float64, float64)]
+    >>> Annotate(add, signature=add.signatures[1])  # annotation for (float64, float64)
     """
-    def __init__(self, function, **kwargs):
+    def __init__(self, function, signature=None, **kwargs):
 
         style = kwargs.get('style', 'default')
         if not function.signatures:
             raise ValueError('function need to be jitted for at least one signature')
-        for sig in function.signatures:
-            ann = function.get_annotation_info(sig)
+        ann = function.get_annotation_info(signature=signature)
         self.ann = ann
 
         for k,v in ann.items():


Test patch:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -12,6 +12,8 @@
 import inspect
 import pickle
 import weakref
+import collections
+from itertools import chain
 
 try:
     import jinja2
@@ -859,6 +861,31 @@ def foo(a, b):
         # Exercise the method
         foo.inspect_types(utils.StringIO())
 
+        # Test output
+        expected = str(foo.overloads[foo.signatures[0]].type_annotation)
+        with captured_stdout() as out:
+            foo.inspect_types()
+        assert expected in out.getvalue()
+
+    def test_inspect_types_with_signature(self):
+        @jit
+        def foo(a):
+            return a + 1
+
+        foo(1)
+        foo(1.0)
+        # Inspect all signatures
+        with captured_stdout() as total:
+            foo.inspect_types()
+        # Inspect first signature
+        with captured_stdout() as first:
+            foo.inspect_types(signature=foo.signatures[0])
+        # Inspect second signature
+        with captured_stdout() as second:
+            foo.inspect_types(signature=foo.signatures[1])
+
+        self.assertEqual(total.getvalue(), first.getvalue() + second.getvalue())
+
     @unittest.skipIf(jinja2 is None, "please install the 'jinja2' package")
     @unittest.skipIf(pygments is None, "please install the 'pygments' package")
     def test_inspect_types_pretty(self):
@@ -887,6 +914,19 @@ def foo(a, b):
         self.assertIn("`file` must be None if `pretty=True`",
                       str(raises.exception))
 
+    def test_get_annotation_info(self):
+        @jit
+        def foo(a):
+            return a + 1
+
+        foo(1)
+        foo(1.3)
+
+        expected = dict(chain.from_iterable(foo.get_annotation_info(i).items()
+                                            for i in foo.signatures))
+        result = foo.get_annotation_info()
+        self.assertEqual(expected, result)
+
     def test_issue_with_array_layout_conflict(self):
         """
         This test an issue with the dispatcher when an array that is both