docments.py

"""Document parameters using comments.

Docs: https://fastcore.fast.ai/docments.html.md"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_docments.ipynb.

# %% auto #0
__all__ = ['empty', 'docstring', 'parse_docstring', 'isdataclass', 'get_dataclass_source', 'get_source', 'get_name', 'qual_name',
           'ann_parts', 'docments', 'sig_source', 'extract_docstrings', 'DocmentTbl', 'DocmentList', 'DocmentText',
           'sig2str', 'can_render', 'ShowDocRenderer', 'MarkdownRenderer']

# %% ../nbs/04_docments.ipynb #4c662acc
import re,ast,inspect
from tokenize import tokenize,COMMENT
from ast import parse,FunctionDef,AsyncFunctionDef,AnnAssign
from io import BytesIO
from textwrap import dedent
from types import SimpleNamespace
from inspect import getsource,isfunction,ismethod,isclass,signature,Parameter,Signature
from dataclasses import dataclass, is_dataclass
from .utils import *
from .meta import delegates
from . import docscrape
from textwrap import fill
from inspect import isclass,getdoc,signature

# %% ../nbs/04_docments.ipynb #53e21187
def docstring(sym):
    "Get docstring for `sym` for functions ad classes"
    if isinstance(sym, str): return sym
    res = getdoc(sym)
    if not res and isclass(sym): res = getdoc(sym.__init__)
    if not res and callable(sym) and not isclass(sym): res = getdoc(sym.__call__)
    return res or ""

# %% ../nbs/04_docments.ipynb #1e0cf854
def parse_docstring(sym):
    "Parse a numpy-style docstring in `sym`"
    return AttrDict(**docscrape.NumpyDocString(docstring(sym)))

# %% ../nbs/04_docments.ipynb #b1a586d1
def isdataclass(s):
    "Check if `s` is a dataclass but not a dataclass' instance"
    return is_dataclass(s) and isclass(s)

# %% ../nbs/04_docments.ipynb #6549c4b2
def get_dataclass_source(s):
    "Get source code for dataclass `s`"
    return getsource(s) if not getattr(s, "__module__") == '__main__' else ""

# %% ../nbs/04_docments.ipynb #cac59989
def get_source(s):
    "Get source code for string, function object or dataclass `s`"
    if isinstance(s,str): return s
    return getsource(s) if isfunction(s) or ismethod(s) else get_dataclass_source(s) if isdataclass(s) else None

# %% ../nbs/04_docments.ipynb #91c0d15f
def _parses(s):
    "Parse source of function, method, or dataclass `s`"
    try: return parse(dedent(get_source(s) or ''))
    except SyntaxError:
        if isinstance(s,str): raise
        return parse('')   # live object whose file changed on disk since import: degrade to no docments

def _tokens(s):
    "Tokenize Python code in string or function object `s`"
    s = get_source(s)
    if not s: return []
    return tokenize(BytesIO(s.encode('utf-8')).readline)

_clean_re = re.compile(r'^\s*#(.*)\s*$')
def _clean_comment(s):
    res = _clean_re.findall(s)
    return res[0] if res else None

def _param_locs(s, returns=True, args_kwargs=False):
    "`dict` of parameter line numbers to names"
    body = _parses(s).body
    if len(body)==1:
        defn = body[0]
        if isinstance(defn, (FunctionDef, AsyncFunctionDef)):
            res = {arg.lineno:arg.arg for arg in defn.args.args}
            # Add *args if present
            if defn.args.vararg: res[defn.args.vararg.lineno] = defn.args.vararg.arg
            # Add keyword-only args
            res.update({arg.lineno:arg.arg for arg in defn.args.kwonlyargs})
            # Add **kwargs if present
            if defn.args.kwarg and args_kwargs: res[defn.args.kwarg.lineno] = defn.args.kwarg.arg
            if returns and defn.returns: res[defn.returns.lineno] = 'return'
            return res
        elif isdataclass(s):
            res = {arg.lineno:arg.target.id for arg in defn.body if isinstance(arg, AnnAssign)}
            return res
    return None

# %% ../nbs/04_docments.ipynb #93930c1f
empty = Parameter.empty

# %% ../nbs/04_docments.ipynb #6434e80e
def _get_comment(line, arg, comments, parms):
    if line in comments: return comments[line].strip()
    line -= 1
    res = []
    while line and line in comments and line not in parms:
        res.append(comments[line])
        line -= 1
    return dedent('\n'.join(reversed(res))) if res else None

def _get_full(p, docs, eval_str=False):
    anno = p.annotation
    if anno==empty:
        if p.default!=empty: anno = type(p.default)
        elif eval_str: anno = None
    return AttrDict(docment=docs.get(p.name), anno=anno, default=p.default, kind=p.kind)

# %% ../nbs/04_docments.ipynb #1b4d817c
def _merge_doc(dm, npdoc):
    if not npdoc: return dm
    if not isinstance(dm, dict): return dm or '\n'.join(npdoc.desc)
    # if not dm.anno or dm.anno==empty: dm.anno = npdoc.type
    if not dm.docment: dm.docment = '\n'.join(npdoc.desc)
    return dm

def _merge_docs(dms, npdocs):
    npparams = npdocs['Parameters']
    params = {nm:_merge_doc(dm,npparams.get(nm,None)) for nm,dm in dms.items()}
    if 'return' in dms: params['return'] = _merge_doc(dms['return'], npdocs['Returns'])
    return params

# %% ../nbs/04_docments.ipynb #40e3f274
def _get_property_name(p):
    "Get the name of property `p`"
    if hasattr(p, 'fget'):
        return p.fget.func.__qualname__ if hasattr(p.fget, 'func') else p.fget.__qualname__
    else: return next(iter(re.findall(r'\'(.*)\'', str(p)))).split('.')[-1]

# %% ../nbs/04_docments.ipynb #da0465e3
def get_name(obj):
    "Get the name of `obj`"
    if isinstance(obj, partial):
        nm = get_name(obj.func)
        args = [repr(a) for a in obj.args] + [f'{k}={repr(v)}' for k,v in obj.keywords.items()]
        return f"{nm}[partial: {', '.join(args)}]"
    if hasattr(obj, '__name__'):       return obj.__name__
    elif getattr(obj, '_name', False): return obj._name
    elif hasattr(obj,'__origin__'):    return str(obj.__origin__).split('.')[-1]
    elif type(obj)==property:          return _get_property_name(obj)
    elif callable(obj):                return type(obj).__name__
    else:                              return str(obj).split('.')[-1]


# %% ../nbs/04_docments.ipynb #5e273f22
def qual_name(obj):
    "Get the qualified name of `obj`"
    if hasattr(obj,'__qualname__'): return obj.__qualname__
    if ismethod(obj):       return f"{get_name(obj.__self__)}.{get_name(fn)}"
    return get_name(obj)

# %% ../nbs/04_docments.ipynb #2e865627
def ann_parts(anno):
    "The underlying type and metadata tuple of an `Annotated`, else `(anno, ())`"
    if typing.get_origin(anno) is not typing.Annotated: return anno, ()
    args = typing.get_args(anno)
    t = None if args[0] is type(None) else args[0]
    return t, args[1:]

# %% ../nbs/04_docments.ipynb #9b62ab20
def docments(s, full=False, eval_str=False, returns=True, args_kwargs=False):
    "Get docments for `s`"
    if isclass(s) and not is_dataclass(s): s = s.__init__
    try: sig = signature_ex(s, eval_str=eval_str)
    except (ValueError, TypeError): return AttrDict()
    nps = parse_docstring(s)
    docs = {}
    while s:
        p = _param_locs(s, returns=returns, args_kwargs=args_kwargs) or {}
        c = {o.start[0]:_clean_comment(o.string) for o in _tokens(s) if o.type==COMMENT}
        for k,v in p.items():
            if v not in docs: docs[v] = _get_comment(k, v, c, p)
        s = getattr(s, '__delwrap__', None)
    
    res = {k:_get_full(v, docs, eval_str=eval_str) if full else docs.get(k) for k,v in sig.parameters.items()}
    if returns:
        if full: res['return'] = AttrDict(docment=docs.get('return'), anno=sig.return_annotation, default=empty)
        else: res['return'] = docs.get('return')
    for k,p in sig.parameters.items():
        if (res[k].docment if full else res[k]) is None:
            am = first(o for o in ann_parts(p.annotation)[1] if isinstance(o,str))
            if am is None: continue
            if full: res[k].docment = am
            else: res[k] = am
    return AttrDict(_merge_docs(res, nps))

# %% ../nbs/04_docments.ipynb #40cdbeb2
def sig_source(obj):
    "Full source of signature line(s) for a function or class."
    src = inspect.getsource(obj)
    tree = ast.parse(src)
    body_start = tree.body[0].body[0].lineno
    if body_start == 1: return src.splitlines()[0]
    return '\n'.join(src.splitlines()[:body_start-1])

# %% ../nbs/04_docments.ipynb #6f0d6ea4
def _get_params(node):
    params = [a.arg for a in node.args.args]
    if node.args.vararg: params.append(f"*{node.args.vararg.arg}")
    if node.args.kwarg: params.append(f"**{node.args.kwarg.arg}")
    return ", ".join(params)

# %% ../nbs/04_docments.ipynb #180b4c2d
class _DocstringExtractor(ast.NodeVisitor):
    def __init__(self): self.docstrings,self.cls,self.cls_init = {},None,None

    def visit_FunctionDef(self, node):
        name = node.name
        if name == '__init__':
            self.cls_init = node
            return
        elif name.startswith('_'): return
        elif self.cls: name = f"{self.cls}.{node.name}"
        docs = ast.get_docstring(node)
        params = _get_params(node)
        if docs: self.docstrings[name] = (docs, params)
        self.generic_visit(node)

    def visit_ClassDef(self, node):
        self.cls,self.cls_init = node.name,None
        docs = ast.get_docstring(node)
        if docs: self.docstrings[node.name] = ()
        self.generic_visit(node)
        if not docs and self.cls_init: docs = ast.get_docstring(self.cls_init)
        params = _get_params(self.cls_init) if self.cls_init else ""
        if docs: self.docstrings[node.name] = (docs, params)
        self.cls,self.cls_init = None,None

    def visit_Module(self, node):
        module_doc = ast.get_docstring(node)
        if module_doc: self.docstrings['_module'] = (module_doc, "")
        self.generic_visit(node)

# %% ../nbs/04_docments.ipynb #b1d612e9
def extract_docstrings(code):
    "Create a dict from function/class/method names to tuples of docstrings and param lists"
    extractor = _DocstringExtractor()
    extractor.visit(ast.parse(code))
    return extractor.docstrings

# %% ../nbs/04_docments.ipynb #e4e1b815
def _non_empty_keys(d:dict): return L([k for k,v in d.items() if v != inspect._empty])
def _bold(s): return f'**{s}**' if s.strip() else s

# %% ../nbs/04_docments.ipynb #ac070254
def _escape_markdown(s):
    for c in '|^': s = re.sub(rf'\\?\{c}', rf'\{c}', s)
    return s.replace('\n', '<br>')

# %% ../nbs/04_docments.ipynb #ced78f56
def _maybe_nm(o):
    if (o == inspect._empty): return ''
    else: return o.__name__ if hasattr(o, '__name__') else str(o)

# %% ../nbs/04_docments.ipynb #fe6d83f1
def _list2row(l:list): return '| '+' | '.join([_maybe_nm(o) for o in l]) + ' |'

# %% ../nbs/04_docments.ipynb #44ac2e4f
class _DocmentBase:
    def __init__(self, obj):
        self.obj,self.dm = obj, docments(obj, full=True)
        if 'self' in self.dm: del self.dm['self']
    
    @property
    def has_docment(self): return any(v.get('docment') for v in self.dm.values())

# %% ../nbs/04_docments.ipynb #b45f731d
class DocmentTbl(_DocmentBase):
    _map = {'anno':'Type', 'default':'Default', 'docment':'Details'}

    def __init__(self, obj, verbose=True, returns=True):
        "Compute the docment table string"
        super().__init__(obj)
        self.verbose = verbose
        self.returns = False if isdataclass(obj) else returns
        try: self.params = L(signature(obj, eval_str=True).parameters.keys())
        except (ValueError,TypeError): self.params=[]
        for d in self.dm.values(): d['docment'] = ifnone(d['docment'], inspect._empty)

    @property
    def _columns(self):
        "Compute the set of fields that have at least one non-empty value so we don't show tables empty columns"
        cols = set(flatten(L(self.dm.values()).filter().map(_non_empty_keys)))
        candidates = self._map if self.verbose else {'docment': 'Details'}
        return {k:v for k,v in candidates.items() if k in cols}

    @property
    def has_docment(self): return 'docment' in self._columns and self._row_list

    @property
    def has_return(self): return self.returns and bool(_non_empty_keys(self.dm.get('return', {})))

    def _row(self, nm, props): return [nm] + [props[c] for c in self._columns]

    @property
    def _row_list(self):
        ordered_params = [(p, self.dm[p]) for p in self.params if p != 'self' and p in self.dm]
        return L([self._row(nm, props) for nm,props in ordered_params])

    @property
    def _hdr_list(self): return ['  '] + [_bold(l) for l in L(self._columns.values())]

    @property
    def hdr_str(self):
        md = _list2row(self._hdr_list)
        return md + '\n' + _list2row(['-' * len(l) for l in self._hdr_list])

    @property
    def params_str(self): return '\n'.join(self._row_list.map(_list2row))

    @property
    def return_str(self): return _list2row(['**Returns**']+[_bold(_maybe_nm(self.dm['return'][c])) for c in self._columns])

    def _repr_markdown_(self):
        if not self.has_docment: return ''
        _tbl = [self.hdr_str, self.params_str]
        if self.has_return: _tbl.append(self.return_str)
        return '\n'.join(_tbl)

    def __eq__(self,other): return self.__str__() == str(other).strip()
    __str__ = _repr_markdown_
    __repr__ = basic_repr()

# %% ../nbs/04_docments.ipynb #5af61ab1
class DocmentList(_DocmentBase):
    def _fmt(self, nm, p):
        anno,default,doc = _maybe_nm(p.get('anno','')), p.get('default',empty), p.get('docment','')
        s = f'`{nm}{":" + anno if anno else ""}{"=" + _maybe_nm(default) if default != empty else ""}`'
        br = '\xa0'*3
        return f'- {s}' + (f'{br}*{doc}*' if doc else '')

    def _repr_markdown_(self): return '\n'.join(self._fmt(k,v) for k,v in self.dm.items())
    __repr__=__str__=_repr_markdown_

# %% ../nbs/04_docments.ipynb #6355b569
def _clean_text_sig(obj):
    if not (sig := getattr(obj, '__text_signature__', None)): return None
    sig = re.sub(r'\$\w+,?\s*', '', sig)
    return get_name(obj) + sig.replace('<unrepresentable>', '...')

# %% ../nbs/04_docments.ipynb #cfcc46bf
def _fmt_sig(name, params, ret_str, maxline, prefix='def'):
    "Format function signature with params and docment comments"
    lines,curr = [],[]
    for fmt,doc in params:
        comment = ' # ' + '\n    # '.join(doc.splitlines()) if doc else ''
        if curr and len(', '.join(curr))+len(fmt)+len(comment.split('\n',1)[0])>maxline:
            lines.append(', '.join(curr) + ',')
            curr = []
        curr.append(fmt)
        if doc: lines.append(', '.join(curr) + ',' + comment); curr = []
    if curr: lines.append(', '.join(curr))
    pstr = '\n    '.join(lines)
    if not pstr: return f"{prefix} {name}({ret_str}"
    return f"{prefix} {name}(\n    {pstr}\n{ret_str}"

# %% ../nbs/04_docments.ipynb #4550820f
def _fmt_default(o):
    if o is empty: return ''
    return o.__name__ if hasattr(o, '__name__') else repr(o)

# %% ../nbs/04_docments.ipynb #7f5e5282
class DocmentText(_DocmentBase):
    def __init__(self, obj, maxline=110, docstring=True):
        super().__init__(obj)
        self.maxline,self.docstring = maxline,docstring

    def _fmt_param(self, nm, p):
        anno,default,kind = p.get('anno',empty), p.get('default',empty), p.get('kind')
        pre = '*' if kind==Parameter.VAR_POSITIONAL else '**' if kind==Parameter.VAR_KEYWORD else ''
        return pre + nm + (f':{_maybe_nm(anno)}' if anno != empty else '') + (f'={_fmt_default(default)}' if default != empty else '')
    
    @property
    def _ret_str(self):
        ret = self.dm.get('return', {})
        anno = f"->{_maybe_nm(ret.get('anno',empty))}" if ret.get('anno',empty) != empty else ''
        doc = f" # {ret['docment']}" if ret.get('docment') else ''
        return f"){anno}:{doc}"
    
    @property
    def params(self): return [(self._fmt_param(k,v), v.get('docment','')) for k,v in self.dm.items() if k != 'return']

    def __str__(self):
        o = self.obj
        if not callable(o):
            doc = docstring(o)
            docstr = f'\n    "{doc}"' if self.docstring and doc else ''
            return f"{type(o).__name__} instance: {truncstr(repr(o), 80)}{docstr}"
        is_inst = callable(o) and not (isfunction(o) or isclass(o) or inspect.isbuiltin(o) or ismethod(o))
        prefix = 'async def' if inspect.iscoroutinefunction(o) else 'def'
        nm = get_name(o) if hasattr(o, '__name__') else f'{type(o).__name__}.__call__' if is_inst else get_name(o)
        if (sig := _clean_text_sig(o)) and not self.params: sig_str = f"{prefix} {sig}"
        else: sig_str = _fmt_sig(nm, self.params, self._ret_str, self.maxline, prefix=prefix)
        doc = getattr(o.__call__, '__doc__', None) if is_inst else o.__doc__
        docstr = f'    "{doc}"' if self.docstring and doc else ''
        return f"{sig_str}\n{docstr}" if docstr else sig_str
    
    __repr__ = __str__
    def _repr_markdown_(self): return f"```python\n{self}\n```"

# %% ../nbs/04_docments.ipynb #d44a7fe3
def sig2str(func, maxline=110):
    "Generate function signature with docments as comments"
    return DocmentText(func, maxline=maxline, docstring=False)

# %% ../nbs/04_docments.ipynb #07287425
def _docstring(sym):
    npdoc = parse_docstring(sym)
    return '\n\n'.join([npdoc['Summary'], npdoc['Extended']]).strip()

def _unwrap_sym(sym):
    if not hasattr(sym, '__signature__'): sym = getattr(sym, '__wrapped__', sym)
    return getattr(sym, 'fget', None) or getattr(sym, 'fset', None) or sym

def can_render(sym):
    "Check if `sym` has a renderable signature"
    sym = _unwrap_sym(sym)
    try: signature(sym, eval_str=True); return True
    except (ValueError, TypeError): return False

# %% ../nbs/04_docments.ipynb #9de89cb6
def _fullname(o):
    module,name = getattr(o, "__module__", None),qual_name(o)
    return name if module is None or module in ('__main__','builtins') else module + '.' + name

class ShowDocRenderer:
    def __init__(self, sym, name:str|None=None, title_level:int=3, maxline:int=110):
        "Show documentation for `sym`"
        sym = _unwrap_sym(sym)
        store_attr()
        self.nm = name or qual_name(sym)
        self.isfunc = inspect.isfunction(sym)
        try: self.sig = signature(sym, eval_str=True)
        except (ValueError,TypeError): self.sig = None
        self.docs = _docstring(sym)
        self.dm = DocmentText(sym, maxline=maxline, docstring=False)
        self.fn = _fullname(sym)

    __repr__ = basic_repr()

# %% ../nbs/04_docments.ipynb #e569d885
def _f_name(o): return f'<function {o.__name__}>' if isinstance(o, FunctionType) else None
def _fmt_anno(o): return inspect.formatannotation(o).strip("'").replace(' ','')

def _show_param(param):
    "Like `Parameter.__str__` except removes: quotes in annos, spaces, ids in reprs"
    kind,res,anno,default = param.kind,param._name,param._annotation,param._default
    kind = '*' if kind==inspect._VAR_POSITIONAL else '**' if kind==inspect._VAR_KEYWORD else ''
    res = kind+res
    if anno is not inspect._empty: res += f':{_f_name(anno) or _fmt_anno(anno)}'
    if default is not inspect._empty: res += f'={_f_name(default) or repr(default)}'
    return res

# %% ../nbs/04_docments.ipynb #59797eb7
def _ital_first(s:str):
    "Surround first line with * for markdown italics, preserving leading spaces"
    return re.sub(r'^(\s*)(.+)', r'\1*\2*', s, count=1)

# %% ../nbs/04_docments.ipynb #9184b175
class MarkdownRenderer(ShowDocRenderer):
    "Markdown renderer for `show_doc`"
    def _repr_markdown_(self):
        doc = f'```python\n{self.dm}\n```'
        if self.docs: doc += f"\n\n{_ital_first(self.docs)}"
        return doc

    def __repr__(self):
        doc = str(self.dm)
        if self.docs: doc += f'"""{self.docs}"""'
        return doc
    __str__=__repr__