Airflow API Reference generator

gistfile1.py

import inspect
import pkgutil
from importlib import import_module

from airflow.hooks.base_hook import BaseHook
from airflow.models import BaseOperator


def find_clazzes(directory, base_class):
    found_classes = set()
    for module_finder, name, ispkg in pkgutil.iter_modules([directory]):
        if ispkg:
            continue

        package_name = module_finder.path.replace('/', '.')
        full_module_name = package_name + '.' +name
        try:
            mod = import_module(full_module_name)
        except ModuleNotFoundError:
            print("Error: ", full_module_name)
            continue
        clazzes = inspect.getmembers(mod, inspect.isclass)
        operators_clazzes = [clazz
                             for name, clazz in clazzes
                             if issubclass(clazz, base_class) and
                             clazz.__module__.startswith(package_name)]
        # print(name)
        for clazz in operators_clazzes:
            found_classes.add("%s.%s" % (clazz.__module__,  clazz.__name__))

    return found_classes


def scan(directory, base_class):

    r = find_clazzes(directory, base_class)

    def by_class_name(clazz): return clazz.rsplit('.', 1)[::-1]

    for clazz in sorted(r, key=by_class_name):
        print('.. autoclass::', clazz)


print('# operators')
print('## airflow/operators')
scan('airflow/operators', BaseOperator)
print('## airflow/contrib/operators')
scan('airflow/contrib/operators', BaseOperator)

print('# sensors')
print('## airflow/sensors')
scan('airflow/sensors', BaseOperator)
print('## airflow/contrib/sensors')
scan('airflow/contrib/sensors', BaseOperator)

print('# hooks')
print('## airflow/hooks')
scan('airflow/hooks', BaseHook)
print('## airflow/contrib/hooks')
scan('airflow/contrib/hooks', BaseHook)