ndoc.py

import numpy.random as nr
import docutils
import docutils.frontend
import docutils.parsers
import docutils.parsers.rst
import numpydoc.docscrape as nd
import pygments
import textwrap
from pygments.lexers import Python3Lexer
from pygments.token import Token
import numpy as np 
import numpy.polynomial.legendre
import matplotlib
import matplotlib.pyplot
from docutils.parsers.rst import directives


from pygments import highlight
from pygments.lexers import PythonLexer, RstLexer 
from pygments.formatters import HtmlFormatter, Terminal256Formatter

class UnknownNode(Exception): pass

import flatlatex
class TokenType:
    
    def __init__(self, stack):
        self.stack = stack
        self._value = None
        
    def __call__(self, value):
        self._value = value
        return self
        
    def __enter__(self):
        if self._value is None:
            raise RuntimeError("did you pass a parameter ?")
        self.stack.append(self._value)
        self._value = None
        
    def __exit__(self, *args, **kwargs):
        self.stack.pop()
        
foot = None
class Visitor:
    
    def __init__(self, root_node):
        self.root_node = root_node
        self._stack = []
        self.token = TokenType(self._stack)

    @property
    def stack(self):
        return [x for x in self._stack]
        
    
    def visit_title_reference(self, title_ref):
        with self.token('Keyword'):
            yield from self.visit(title_ref)
    
    def visit_unknown(self, node):
        with self.token('Unknown.'+node.__class__.__name__):
            yield from self.visit(node.children)
        # raise UnknownNode(node.__class__.__name__, list(self.visit(node)))
        
    def visit_literal_block(self, literal_block):
        yield from self.visit(literal_block)
        
    def visit(self, nodes):
        for node in nodes: 
            node_name = node.__class__.__name__
            visitor = getattr(self, 'visit_'+node_name, self.visit_unknown)
            yield from visitor(node)
    
    def visit_document(self, document):
        yield from self.visit(document)
        
    def visit_doctest_block(self, doctest):
        assert len(doctest.children) == 1
        #assert doctest.children[0].__class__.__name__ == 'text', doctest.children
        for a,b,c in Python3Lexer().get_tokens_unprocessed(doctest.children[0].astext()):
            yield b, c
        yield ['Text'], '\n\n'


        
    def visit_literal(self, literal):
        with self.token('Literal'):
            yield from self.visit(literal)
            
    def visit_math(self, math):
        with self.token('Keyword'):
            converter = flatlatex.converter()
            assert len(math.children) == 1
            yield ['Keyword'], converter.convert(math.children[0].astext())
            
    def visit_math_block(self, math):
        with self.token('Keyword'):
            converter = flatlatex.converter()
            assert len(math.children) == 1
            yield ['Text'], '    '
            yield ['Keyword'], converter.convert(math.children[0].astext())
            yield ['Text'], '\n\n'
            
    def visit_paragraph(self, paragraph):
        yield from self.visit(paragraph)
        yield ['Text'], '\n\n'
        
    def visit_section(self, section):
        yield from self.visit(section)
        
    def visit_label(self, label):
        # do noting : footnote
        return []
    def visit_footnote(self, footnote):
        try:
            paragraph = footnote.children[1]
            reference = paragraph.children[1][0]
            yield ['Text'] , paragraph[0].astext().rstrip('\n, ')
            yield ['Text'] , ' : '
            yield ['Text'] , reference.astext()
        except:
            pass


    def visit_system_message(self, message):
        with self.token('SystemMessage'):
            return
            yield ['Mesage'] , "\n<========\n"
            yield from self.visit(message)
            yield ['Mesage'] , "\n========>\n"
        
    def visit_Text(self, text):
        yield (self.stack, text.astext())    
        
        
    def traverse(self):
        yield from self.visit([self.root_node])
        


def fixed_token_generator():
    n = 0 
    type_, text = yield
    while True:
        type_, text = yield n, type_,text
        n += len(text)

def fix_tokens(tokens):
    n = 0
    for type_,text in tokens:
        yield n,type_,text
        n += len(text)
        
def indent_gen(iterator):
     
    
    for a,b in iterator[:-1]:
        yield a,b.replace('\n', '\n    ')
    yield iterator[-1]
    
        

class FakeLexer(pygments.lexer.Lexer):
    
    def get_tokens_unprocessed(self, text):
        yield from fix_tokens(self._get_tokens_unprocessed(text))
    
    def _gen_heading(self, name):
        yield Token.Generic.Heading, name
        #yield Token.Generic.Heading, '\n'
        #yield Token.Generic.Heading, '='*len(name)
        yield Token.Generic.Heading, '\n'
 
    
    def _gen_rst_section(self, name, sections):
        section = sections[name]
        if not section:
            return
        yield from self._gen_heading(name)
        yield Token.Text, '\n'
        yield from self.tokens_from_sections('\n'.join(section))
        
    def _gen_typed_parameter_section(self, name, sections):
        section = sections[name]
        if not section:
            return
        yield from self._gen_heading(name)
        yield Token.Text, '\n'
        for name, type_, desc in section:
            yield Token.Keyword, name
            yield Token.Text, ' : '
            yield Token.Text, type_
            yield Token.Text, '\n'
            yield Token.Text, ' '*4
            yield from indent_gen(self.tokens_from_sections('\n'.join(desc))) # need to be 4 indented.
    
    def _get_tokens_unprocessed(self, text):
        sections = nd.NumpyDocString(text)
        notes = sections['Notes']
        yield Token.Text, sections['Signature']
        yield Token.Text, '\n\n'
        yield Token.Text, '\n'.join(sections['Summary'])
        yield Token.Text, '\n\n'
        yield from self.tokens_from_sections('\n'.join(sections['Extended Summary']))
        
        yield Token.Text, '\n'
        
        
        yield from self._gen_typed_parameter_section('Parameters', sections)
        yield from self._gen_typed_parameter_section('Returns', sections)
        yield from self._gen_typed_parameter_section('Yields', sections)
        
        yield from self._gen_typed_parameter_section('Raises', sections)
        yield from self._gen_typed_parameter_section('Warns', sections)
        yield from self._gen_rst_section('Other Parameters', sections)
        yield from self._gen_rst_section('Attributes', sections)

        name = 'See Also'
        section = sections[name]
        
        
        if section:
            # TODO yield proper token here.
            yield from self._gen_heading(name)
            yield Token.Text, '\n'
            nodesc = [a for (a,b,c) in section if (not b and not c)]
            text =', '.join(nodesc)
            yield from self.tokens_from_sections(text)
            
            withdesc = [(a,b,c) for a,b,c in section if (b or c)]
            text ='\n'.join(['%s : %s' % (x1 , ''.join(x2)) for x1,x2,x3 in withdesc])
            yield from self.tokens_from_sections(text)


        yield from self._gen_rst_section('Notes', sections)
        yield from self._gen_rst_section('Warnings', sections)
        yield from self._gen_rst_section('References', sections)
        yield from self._gen_rst_section('Examples', sections)
    

    
    def tokens_from_sections(self, text):
        
        default_settings = docutils.frontend.OptionParser(
            components=(docutils.parsers.rst.Parser,)
        ).get_default_values()
        document = docutils.utils.new_document('fileobj.name', default_settings)
        parser = docutils.parsers.rst.Parser()
        p = parser.parse(text, document)
        
        items = Visitor(document).traverse()
        
        ll = []
        for item in  items:
            tt = Token
            for att in item[0]:
                tt = getattr(tt, att)
            ll.append((tt , item[1]))
        return ll
        return list(Visitor(document).traverse())
        #visitor = LinkCheckerVisitor(document)
        #document.walkabout(visitor)
        #if visitor.aborted:
        #    raise ValueError('Unknown thing', visitor.aborted)
        #return visitor._tokens

def ndoc(meth):
    default_settings = docutils.frontend.OptionParser(
                components=(docutils.parsers.rst.Parser,)
            ).get_default_values()
    document = docutils.utils.new_document('fileobj.name', default_settings)

    # meth = nr.multivariate_normal
    # meth = nr.RandomState.choice
    # meth = np.polynomial.legendre.legfit

    # meth = matplotlib.pyplot.plot
    code = meth.__doc__
    try:
        print(highlight(code, FakeLexer(), Terminal256Formatter()))
    except ValueError:
        print(code)