collect.py

import taichi as ti
from types import ModuleType, FunctionType
import inspect
import pdb
import csv

# FIXME: Without this init some calls will fail
ti.init()

def has_doc_string(x):
    return getattr(x, '__doc__', None) is not None

def get_name(x):
    return x.__name__ if getattr(x, '__name__', None) is not None else x

queue = []
# Dict{level : List[Tuple(api: object, api_name: str)]}
res = {0: [], 1: []}
def search():
    queue.append((ti, 0, 'ti'))

    while queue:
        s, level, prefix = queue.pop(0)
        if level == 2:
            break
        if s == ti.lib.taichi_core:
            continue

        for x in dir(s):
            if x.startswith('_'):
                continue
            api = getattr(s, x)
            res[level].append((api, f'{prefix}.{x}'))
            if isinstance(api, ModuleType) and api.__name__.startswith('taichi.'):
                queue.append((api, level+1, f'{prefix}.{x}'))


def write_to_csv(filename, api_list):
    with open(f'{filename}.csv', 'w') as f:
        writer = csv.writer(f)
        writer.writerow([filename])

        for api_name in api_list:
            if isinstance(api_name, str):
                writer.writerow([api_name])
            else:
                pdb.set_trace()
                writer.writerow([api_name])

def main():
    search()
    top_level = res[0]
    second_level = res[1]

    # 1. Do we really want to expose them here? Consider:
    #    - remove from top level
    #    - don't expose
    top_level_apis = [name for api, name in top_level]
    write_to_csv('top_level_apis', top_level_apis)

    # 2. Top level apis must have doc strings!
    missing_doc_string = [name for api, name in top_level if not has_doc_string(api)]
    write_to_csv('missing_doc_string', missing_doc_string)


    # 3. Do we really want to expose these modules? Consider:
    #    - Don't expose, e.g. taichi.lib, taichi.lang.type_factory_impl
    #    - Rename, short and representable, e.g. taichi.torch_io -> taichi.io.from/to_torch
    #    Ignore imported python modules and taichi module itself
    top_level_modules = [name for api, name in top_level if isinstance(api, ModuleType) and api.__name__.startswith('taichi.')]
    write_to_csv('top_level_modules', top_level_modules)



if __name__ == '__main__':
    main()