Creating Type Stubs for Scientific Python (Part 4)

class Normalizer(Interpreter):
    def configure(self, module:str|None, classes: dict|None, is_param:bool):
        if module is None:
            module = ''
        x = module.find('.')
        if x >= 0:
            self._tlmodule = module[:x]  # top-level module
        else:
            self._tlmodule = module
        self._module = module
        self._classes = classes
        self._is_param = is_param

    def handle_qualname(self, name: str, imports: set) -> str:
        return name
    
    def start(self, tree) -> tuple[str, set[str]|None]:
        result = self.visit(tree.children[0])
        return result
        
    def type_list(self, tree) -> tuple[str, set[str]|None]:
        types = [] # We want to preserve order so don't use a set
        imports = set()
        literals = []
        has_none = False
        for child in tree.children:
            if isinstance(child, Tree):
                result = self._visit_tree(child)
                if result:
                    if result[0] == 'None':
                        has_none = True
                        continue
                    if result[0].startswith('Literal:'):
                        literals.append(result[0][8:])
                    else:
                        type = result[0]
                        if type not in types:
                            types.append(type)
                    if result[1]:
                        imports.update(result[1])

        if not imports:
            imports = None
        if literals:
            type = 'Literal[' + ','.join(literals) + ']'
            if type not in types:
                types.append(type)
        if has_none:
            type = 'None'
            if type not in types:
                types.append(type)
        type = '|'.join(types)
        return type, imports

    def type(self, tree)-> tuple[str, set[str]|None]:
        for child in tree.children:
            if isinstance(child, Tree):
                result = self._visit_tree(child)
                if result:
                    return result
        assert(False)

    _basic_types = {
        'ANY': 'Any',
        'INT': 'int',
        'STR' : 'str',
        'FLOAT': 'float',
        'BOOL': 'bool',
        'SCALAR': 'Scalar',
        'COMPLEX': 'complex',
        'OBJECT': 'Any',
        'PATHLIKE': 'PathLike',
        'FILELIKE': 'FileLike'
    }

    def array_type(self, tree):
        arr_types = set()
        elt_type = None
        imports = set()
        for child in tree.children:
            if isinstance(child, Token) and (child.type == 'NDARRAY' or child.type == 'NUMPY'):
                arr_types.add('NDArray')
                imports.add(('NDArray', 'numpy.typing'))
            if isinstance(child, Tree) and isinstance(child.data, Token):
                tok = child.data
                subrule = tok.value
                if subrule == 'array_kinds':
                    types, imp = self._visit_tree(child)
                    arr_types.update(types)
                    imports.update(imp)
                elif subrule == 'array_kind':
                    type, imp = self._visit_tree(child)
                    arr_types.add(type)
                    imports.update(imp)
                elif subrule == 'basic_type' or subrule == 'type_qualifier':
                    elt_type, imp = self._visit_tree(child)
                    imports.update(imp)
        if elt_type:
            if self._is_param and 'list' in arr_types:
                arr_types.add('Sequence')
                arr_types.remove('list')
                imports.add(('Sequence', 'typing'))
            return '|'.join([f'{typ}[{elt_type}]' for typ in arr_types]), imports
        else:
            return '|'.join(arr_types), imports

    def array_kinds(self, tree):
        imports = set()
        types = set()
        for child in tree.children:
            if isinstance(child, Tree):
                type, imp = self._visit_tree(child)
                imports.update(imp)
                types.add(type)
        return types, imports

    def array_kind(self, tree):
        arr_type = ''
        imports = set()
        for child in tree.children:
            if isinstance(child, Token):
                if child.type == 'NDARRAY':
                    arr_type = 'NDArray'
                    imports.add(('NDArray', 'numpy.typing'))
                elif self._is_param or child.type == 'ARRAYLIKE':
                    arr_type = 'ArrayLike'
                    imports.add(('ArrayLike', 'typing'))
                elif child.type == 'LIST':
                    arr_type = 'list'
                elif child.type == 'SEQUENCE':
                    arr_type = 'Sequence'
                    imports.add(('Sequence', 'typing'))
                else:
                    continue
                break

        if not arr_type:
            arr_type = 'NDArray'
            imports.add(('NDArray', 'numpy.typing'))

        return arr_type, imports

    def type_qualifier(self, tree):
        imports = set()
        for child in tree.children:
            if isinstance(child, Tree):
                return self._visit_tree(child)
            elif isinstance(child, Token):
                if child.type == 'QUALNAME':
                    type = self.handle_qualname(child.value, imports)
                    return type, imports
        # OF ARRAYS falls through here
        imports.add(('ArrayLike', 'numpy.typing'))
        return 'ArrayLike', imports

    def basic_type(self, tree) -> tuple[str, set[str]|None]:
        imports = set()
        for child in tree.children:
            if isinstance(child, Token):
                if child.type == 'ANY':
                    imports.add(('Any', 'typing'))
                    return 'Any', imports
                elif child.type == 'SCALAR':
                    imports.add(('Scalar', f'{self._tlmodule}._typing'))
                    return 'Scalar', imports
                elif child.type == 'PATHLIKE':
                    imports.add(('PathLike', 'os'))
                elif child.type == 'FILELIKE':
                    imports.add(('IO', 'typing'))
                    return 'IO', imports
                if child.type in self._basic_types:
                    return self._basic_types[child.type], imports

        assert(False)

    def callable_type(self, tree):
        # TODO: handle signature
        imports = set()
        imports.add(('Callable', 'typing'))
        return "Callable", imports

    def class_type(self, tree):
        cname = ''
        for child in tree.children:
            if isinstance(child, Tree):
                return self._visit_tree(child)
        assert(False)
        
    def class_specifier(self, tree):
        imp = set()
        cname = ''
        for child in tree.children:
            if isinstance(child, Token) and child.type == 'QUALNAME':
                cname = child.value
        # Now we need to normalize the name and find the imports
        x = cname.rfind('.')
        if x > 0:
            imp.add((cname[x+1:], cname[:x]))
            cname = cname[x+1:]
        elif self._classes and cname in self._classes:
            imp.add((cname, self._classes[cname]))

        return cname, imp

    def dict_type(self, tree):
        dict_type = ''
        from_type = None
        to_type = None
        imports = set()
        for child in tree.children:
            if isinstance(child, Token):
                if child.type == 'MAPPING':
                    dict_type = 'Mapping'
                    imports.add(('Mapping', 'typing'))
                elif child.type == 'DICT':
                    dict_type = 'dict'
                elif child.type == 'QUALNAME':
                    to_type = self.handle_qualname(child.value, imports)
                    if from_type is None:
                        from_type = to_type
            elif isinstance(child, Tree):
                to_type, imp = self._visit_tree(child)
                if imp:
                    imports.update(imp)
                if from_type is None:
                    from_type = to_type

        if from_type is not None:
            dict_type += f'[{from_type}, {to_type}]'

        return dict_type, imports

    
    def qualname(self, tree):
        for child in tree.children:
            if isinstance(child, Token):
                imports = set()
                return self.handle_qualname(child.value, imports), imports
            

    def filelike_type(self, tree):
        imports = set()
        imports.add(('FileLike', f'{self._tlmodule}._typing'))
        return 'FileLike', imports

    def generator_type(self, tree):
        imports = set()
        imports.add(('Generator', 'collections.abc'))
        for child in tree.children:
            if isinstance(child, Tree):
                type, imp = self._visit_tree(child)
                if type:
                    imports.update(imp)
                    return f'Generator[{type}, None, None]', imports
        return 'Generator', imports

    def iterable_type(self, tree):
        imports = set()
        imports.add(('Iterable', 'collections.abc'))
        for child in tree.children:
            if isinstance(child, Tree):
                type, imp = self._visit_tree(child)
                if type:
                    imports.update(imp)
                    return f'Iterable[{type}]', imports
        return 'Iterable', imports
    
    def iterator_type(self, tree):
        imports = set()
        imports.add(('Iterator', 'collections.abc'))
        for child in tree.children:
            if isinstance(child, Tree):
                type, imp = self._visit_tree(child)
                if type:
                    imports.update(imp)
                    return f'Iterator[{type}]', imports
        return 'Iterator', imports
    
    def literal_type(self, tree)-> tuple[str, set[str]|None]:
        imports = set()
        imports.add(('Literal', 'typing'))
        assert(len(tree.children) == 1 and isinstance(tree.children[0], Token))
        tok = tree.children[0]
        type = tok.type
        if type == 'STRING' or type == 'NUMBER':
            return f'Literal:' + tok.value, imports
        if type == 'NONE':
            return 'None', set()
        if type == 'TRUE':
            return 'Literal:True', imports
        if type == 'FALSE':
            return 'Literal:False', imports
        assert(False)

    def optional_type(self, tree):
        for child in tree.children:
            if isinstance(child, Tree):
                type, imports = self._visit_tree(child)
                type += '|None'
                return type, imports
        return 'None', set()

    def restricted_type(self, tree):
        imp = set()
        types = []
        values = []
        rtn = ''
        for child in tree.children:
            if isinstance(child, Tree):
                type, imports = self._visit_tree(child)
                imp.update(imports)
                if type.startswith('Literal:'):
                    values.append(type[8:])
                else:
                    types.append('None')
            elif isinstance(child, Token) and child.type == 'STR':
                types.append('str')
        if values:
            types.append(f'Literal[{",".join(values)}]')
            imp.add(('Literal', 'typing'))
        return '|'.join(types), imp

    def set_type(self, tree):
        set_type = None
        elt_type = None
        imports = set()
        for child in tree.children:
            if isinstance(child, Token):
                if child.type == 'SET':
                    set_type = 'set'
                elif child.type == 'FROZENSET':
                    set_type = 'frozenset'
            elif isinstance(child, Tree):
                elt_type, imports = self._visit_tree(child)

        if elt_type:
            set_type += '[' + elt_type + ']'

        return set_type, imports

    def tuple_type(self, tree):
        types = []
        imp = set()
        repeating = False
        count = 1
        has_shape = False
        for child in tree.children:
            if isinstance(child, Tree):
                type, imports = self._visit_tree(child)
                if type is None:
                    has_shape = True
                else:
                    types.append(type)
                    imp.update(imports)
            elif isinstance(child, Token) and child.type in ['OF', 'WITH']:
                repeating = True
            elif isinstance(child, Token) and child.type == 'NUMBER':
                count = int(child.value)
        if has_shape:
            return 'tuple', None
        if types:
            if repeating:
                if count > 1:
                    types = ["|".join(types)] * count
                else:
                    return f'tuple[{"|".join(types)}, ...]', imp
            
            return f'tuple[{",".join(types)}]', imp
        return 'tuple', None

    def union_type(self, tree):
        types = set()
        imports = set()
        for child in tree.children:
            if isinstance(child, Tree):
                type, imp = self._visit_tree(child)
                types.add(type)
                imports.update(imp)
        return '|'.join(types), imports
    
    def shape(self, tree):
        return None, None
    

    def optional_type(self, tree):
        for child in tree.children:
            if isinstance(child, Tree):
                type, imports = self._visit_tree(child)
                type += '|None'
                return type, imports
        return 'None', set()

_lark = Lark(_grammar)
_norm =  _norm = Normalizer()

    
def parse_type(s: str, modname: str|None = None, classes: dict|None = None, is_param:bool=False) -> tuple[str, dict[str, list[str]]]:
    """ Parse a type description from a docstring, returning the normalized
        type and the set of required imports, or None if no imports are needed.
    """
    try:
        tree = _lark.parse(s)
        _norm.configure(modname, classes, is_param)
        n = _norm.visit(tree)
        # Collect, dedup and sort the imports
        imps = {}
        if n[1]:
            for imp in n[1]:
                what, where = imp
                if where not in imps:
                    imps[where] = []
                imps[where].append(what)
            # Sort the imports for a module
            for where in imps.keys():
                imps[where] = sorted(imps[where])
        return n[0], imps
    except Exception as e:
        return s, {}

def normalize_type(s: str, modname: str|None = None, classes: dict|None = None, is_param: bool = False) -> tuple[str|None, dict[str, list[str]]]:
    try:
        return parse_type(remove_shape(s), modname, classes, is_param)
    except Exception as e:
        return None, {}

_basic_types = {
    # Key: lower() version of type
    'any': 'Any', 
    'array': 'NDArray',
    'arraylike': 'NDArray',
    'array-like': 'NDArray',
    'bool': 'bool',
    'bools': 'bool',
    'boolean': 'bool',
    'booleans': 'bool',
    'bytearray': 'bytearray',
    'callable': 'Callable',
    'complex': 'complex',
    'dict': 'dict',
    'dictionary': 'dict',
    'dictionaries': 'dict',
    'filelike': 'FileLike',
    'file-like': 'FileLike',
    'float': 'float',
    'floats': 'float',
    'frozenset': 'frozenset',
    'int': 'int',
    'ints': 'int',
    'iterable': 'Iterable',
    'list': 'list',
    'memoryview': 'memoryview',
    'ndarray': 'np.ndarray',
    'none': 'None',
    'object': 'Any',
    'objects': 'Any',
    'pathlike': 'PathLike',
    'path-like': 'PathLike',
    'range': 'range',
    'scalar': 'Scalar',
    'sequence': 'Sequence',
    'set': 'set',
    'str': 'str',
    'string': 'str',
    'strings': 'str',
    'tuple': 'tuple',
    'tuples': 'tuple',
}

def is_trivial(s, modname: str, classes: set|dict|None = None):
    """
    Returns true if the docstring is trivially and unambiguously convertible to a 
    type annotation, and thus need not be written to the map file for further
    tweaking.

    s - the type docstring to check
    modname - the module name
    classes - a set of class names or dictionary keyed on classnames 
    """
    s = s.strip()
    sl = remove_shape(s.lower())
    if sl.endswith(" objects"):
        sl = sl[:-8]

    if sl in _basic_types:
        return True

    # Check if it's a string

    if sl and sl[0] == sl[-1] and (sl[0] == '"' or sl[0] =="'"):
        return True

    if classes:
        if s in classes or (_ident.match(s) and s.startswith(modname + '.')):
            return True

    # We have to watch out for ambiguous things like 'list of str or bool'.
    # This is just a kludge to look for both 'of' and 'or' in the type and
    # reject it.
    x1 = sl.find(' of ')
    x2 = sl.find(' or ')
    if x1 >= 0 and x2 > x1:
        return False
    
    # Handle tuples
    if sl.startswith('tuple'):
        sx = s[5:].strip()
        if not sx:
            return True
        if sx[0] in '({[' and sx[-1] in '})]':
            # TODO We should make sure there are no other occurences of these
            # A lot of this is getting to where we should go back to regexps.
            return is_trivial(sx[1:-1], modname, classes)
        
        # Strip off leading OF or WITH
        if sx.startswith ('of ') or sx.startswith('with '):
            x = sx.find(' ')
            sx = sx[x+1:].strip()
        
        # Strip off a number
        if sx and sx[0].isdigit():
            x = sx.find(' ')
            sx = sx[x+1:]

        if is_trivial(sx, modname, classes):
            return True
        
    for s1 in [s, s.replace(',', ' or ')]:
        for splitter in [' or ', '|']:
            if s1.find(splitter) > 0:
                if all([len(c.strip())==0 or is_trivial(c.strip(), modname, classes) \
                        for c in s1.split(splitter)]):
                    return True
        
    if s.find(' of ') > 0:
        # Things like sequence of int, set of str, etc
        parts = s.split(' of ')
        if len(parts) == 2 and is_trivial(parts[0], modname, None) and is_trivial(parts[1], modname, classes):
            return True
        
    # Handle restricted values in {}
    if s.startswith('{') and s.endswith('}'):
        parts = s[1:-1].split(',')
        parts = [p.strip() for p in parts]
        return all([is_trivial(p, modname, None) for p in parts])
    
    return False