1"""Tools for ipython notebooks."""
2
3
4def print_source_in_module(function, module):  # pragma: no cover
5    """
6    For use inside an jupyter_ notebook: given a module and a function, print the source code.
7
8    Based on:
9
10        http://stackoverflow.com/questions/20665118/how-to-show-source-code-of-a-package-function-in-ipython-notebook
11    """
12    from inspect import getmembers, isfunction, getsource
13    from pygments import highlight
14    from pygments.lexers import PythonLexer
15    from pygments.formatters import HtmlFormatter
16    from IPython.core.display import HTML
17
18    internal_module = __import__(module)
19    internal_functions = dict(getmembers(internal_module, isfunction))
20    return HTML(highlight(getsource(internal_functions[function]), PythonLexer(), HtmlFormatter(full=True)))
21
22
23def print_source(function, **kwargs):  # pragma: no cover
24    """
25    For use inside a jupyter_ notebook: given a function, print the source code.
26
27    Args:
28        **kwargs: Passed to HtmlFormatter
29
30    Return:
31        HTML string.
32    """
33    from inspect import getsource
34    from pygments import highlight
35    from pygments.lexers import PythonLexer
36    from pygments.formatters import HtmlFormatter
37    from IPython.core.display import HTML
38
39    if "full" not in kwargs: kwargs["full"] = True
40    return HTML(highlight(getsource(function), PythonLexer(), HtmlFormatter(**kwargs)))
41
42
43def print_doc(function, **kwargs):  # pragma: no cover
44    """
45    For use inside a jupyter_ notebook: given a function, print the docstring.
46
47    Args:
48        **kwargs: Passed to HtmlFormatter
49
50    Return:
51        HTML string.
52    """
53    from inspect import getsource
54    from pygments import highlight
55    from pygments.lexers import PythonLexer
56    from pygments.formatters import HtmlFormatter
57    from IPython.core.display import HTML
58
59    # Extract source code up to end of docstring.
60    lines, count = [], 0
61    for l in getsource(function).splitlines():
62        lines.append(l)
63        if l.lstrip().startswith('"""'): count += 1
64        if count == 2: break
65
66    if "full" not in kwargs: kwargs["full"] = True
67    return HTML(highlight("\n".join(lines), PythonLexer(), HtmlFormatter(**kwargs)))
68
69
70def ipw_listdir(top=".", recurse=True, widget_type="dropdown"):   # pragma: no cover
71    """
72    Return an ipython widget listing all the files located within the directory ``top``
73    that can be inspected with :ref:`abiopen.py`. The user can select the file in the widget
74    and print info on the corresponding file inside the notebook.
75
76    Args:
77        top: Initial directory.
78        recurse: False to ignore directories within ``top``.
79        widget_type: Specify the widget to create. Possible values in:
80            ["tooglebuttons", "dropdown", "radiobuttons"]
81    """
82    from abipy import abilab
83    from IPython.display import display, clear_output
84    import ipywidgets as ipw
85
86    # Select the widget class from widget_type
87    d = dict(
88        tooglebuttons=ipw.ToggleButtons,
89        dropdown=ipw.Dropdown,
90        radiobuttons=ipw.RadioButtons,
91    )
92    try:
93        widget_class = d[widget_type]
94    except KeyError:
95        raise KeyError("Invalid `widget_type`: %s, Choose among: %s" % (widget_type, str(list(d.keys()))))
96
97    def on_value_change(change):
98        """Callback"""
99        clear_output()
100        path = change["new"]
101        #print(change)
102        with abilab.abiopen(path) as abifile:
103            print(abifile)
104            #display(abifile)
105
106    # Get dict: dirname --> list_of_files supported by abiopen.
107    dir2files = abilab.dir2abifiles(top, recurse=recurse)
108    children = []
109    for dirname, files in dir2files.items():
110        w = widget_class(options=files, description="%s:" % dirname)
111        # TODO: Should register the callback of "selected" but I didn't find the event type!
112        w.observe(on_value_change, names='value', type="change")
113        children.append(w)
114    box = ipw.VBox(children=children)
115
116    return display(box)
117