automated_api.py

"""Create public API functions based on ServerAPI methods.

Public functions are created in '_api.py' file and imported in '__init_.py'.
The script reads the 'ServerAPI' class and creates functions with the same
signature and docstring in '_api.py' and '__init__.py' with new/removed
functions.

The script is executed by running 'python automated_api.py' in the terminal.

TODOs:
    Use same signature in api functions as is used in 'ServerAPI' methods.
        Right now is used only '(*args, **kwargs)' signature.
    Prepare CI or pre-commit hook to run the script automatically.
"""

import os
import sys
import re
import inspect
import typing

# Fake modules to avoid import errors
requests = type(sys)("requests")
requests.__dict__["Response"] = type(
    "Response", (), {"__module__": "requests"}
)
websocket = type(sys)("websocket")
websocket.__dict__.update({
    name: type(name, (), {"__module__": "websocket"})
    for name in (
        "WebSocket",
        "ABNF",
        "WebSocketProtocolException",
        "WebSocketConnectionClosedException",
        "WebSocketTimeoutException",
    )
})

sys.modules["requests"] = requests
sys.modules["websocket"] = websocket
sys.modules["unidecode"] = type(sys)("unidecode")

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

EXCLUDED_METHODS = {
    "get_default_service_username",
    "get_default_settings_variant",
    "validate_token",
    "set_token",
    "reset_token",
    "create_session",
    "close_session",
    "as_username",
    "validate_server_availability",
    "get_headers",
    "login",
    "logout",
    "set_default_service_username",
}
EXCLUDED_IMPORT_NAMES = {"GlobalContext"}
AUTOMATED_COMMENT = """
# ------------------------------------------------
#     This content is generated automatically.
# ------------------------------------------------
""".strip()


# Read init file and remove ._api imports
def prepare_init_without_api(init_filepath):
    with open(init_filepath, "r") as stream:
        content = stream.read()

    api_regex = re.compile(r"from \._api import \((?P<functions>[^\)]*)\)")
    api_imports = api_regex.search(content)
    start, end = api_imports.span()
    api_imports_text = content[start:end]
    functions_text = api_imports.group("functions")
    function_names = [
        line.strip().rstrip(",")
        for line in functions_text.split("\n")
        if line.strip()
    ]
    function_names_q = {
        f'"{name}"' for name in function_names
    }

    all_regex = re.compile(r"__all__ = \([^\)]*\)")
    all_content = all_regex.search(content)
    start, end = all_content.span()
    all_content_text = content[start:end]
    filtered_lines = []
    for line in content[start:end].split("\n"):
        found = False
        for name in function_names_q:
            if name in line:
                found = True
                break
        if not found:
            filtered_lines.append(line)
    new_all_content_text = (
        "\n".join(filtered_lines).rstrip(") \n") + "\n\n{all_content}\n)"
    )

    return (
        content
        .replace(api_imports_text, "{api_imports}")
        .replace(all_content_text, new_all_content_text)
    ).rstrip("\n")


# Creation of _api.py content
def indent_lines(src_str, indent=1):
    new_lines = []
    for line in src_str.split("\n"):
        if line:
            line = f"{'    ' * indent}{line}"
        new_lines.append(line)
    return "\n".join(new_lines)


def prepare_docstring(func):
    docstring = inspect.getdoc(func)
    if not docstring:
        return ""

    line_char = ""
    if "\n" in docstring:
        line_char = "\n\n"
    return f'"""{docstring}{line_char}"""'


def _find_obj(obj_full, api_globals):
    parts = list(reversed(obj_full.split(".")))
    _name = None
    for part in parts:
        if _name is None:
            _name = part
        else:
            _name = f"{part}.{_name}"
        try:
            # Test if typehint is valid for known '_api' content
            exec(f"_: {_name} = None", api_globals)
            return _name
        except NameError:
            pass
    return None


def _get_typehint(annotation, api_globals):
    if isinstance(annotation, str):
        annotation = annotation.replace("'", "")

    if inspect.isclass(annotation):
        module_name = str(annotation.__module__)
        full_name = annotation.__name__
        if module_name:
            full_name = f"{module_name}.{full_name}"
        obj_name = _find_obj(full_name, api_globals)
        if obj_name is not None:
            return obj_name

        print("Unknown typehint:", full_name)
        return full_name

    typehint = (
        str(annotation)
        .replace("NoneType", "None")
    )

    full_path_regex = re.compile(
        r"(?P<full>(?P<name>[a-zA-Z0-9_\.]+))"
    )

    for item in full_path_regex.finditer(str(typehint)):
        groups = item.groupdict()
        name = groups["name"]
        obj_name = _find_obj(name, api_globals)
        if obj_name:
            name = obj_name
        else:
            name = name.split(".")[-1]
        typehint = typehint.replace(groups["full"], name)

    forwardref_regex = re.compile(
        r"(?P<full>ForwardRef\('(?P<name>[a-zA-Z0-9]+)'\))"
    )
    for item in forwardref_regex.finditer(str(typehint)):
        groups = item.groupdict()
        name = groups["name"]
        obj_name = _find_obj(name, api_globals)
        if obj_name:
            name = obj_name
        else:
            name = name.split(".")[-1]
        typehint = typehint.replace(groups["full"], name)

    if "Literal" in typehint:
        for match in re.finditer(
            r"(?P<fullcontent>Literal\[(?P<content>[^\]]*)\])", typehint
        ):
            full_content = match.group("fullcontent")
            content = match.group("content")
            items = [f'"{i.strip()}"' for i in content.split(",")]
            new_content = ", ".join(items)
            new_full_content = full_content.replace(content, new_content)
            typehint = typehint.replace(full_content, new_full_content)

    try:
        # Test if typehint is valid for known '_api' content
        exec(f"_: {typehint} = None", api_globals)
        return typehint
    except Exception:
        print("Error while processing typehint:", typehint)
        raise


def _get_param_typehint(param, api_globals):
    if param.annotation is inspect.Parameter.empty:
        return None
    return _get_typehint(param.annotation, api_globals)


def _add_typehint(param_name, param, api_globals):
    typehint = _get_param_typehint(param, api_globals)
    if not typehint:
        return param_name
    return f"{param_name}: {typehint}"


def _kw_default_to_str(param_name, param, api_globals):
    from ayon_api._api_helpers.base import _PLACEHOLDER
    from ayon_api.utils import NOT_SET

    if param.default is inspect.Parameter.empty:
        return _add_typehint(param_name, param, api_globals)

    default = param.default
    if default is _PLACEHOLDER:
        default = "_PLACEHOLDER"
    elif default is NOT_SET:
        default = "NOT_SET"
    elif (
        default is not None
        and not isinstance(default, (str, bool, int, float))
    ):
        raise TypeError("Unknown default value type")
    else:
        default = repr(default)
    typehint = _get_param_typehint(param, api_globals)
    if typehint:
        return f"{param_name}: {typehint} = {default}"
    return f"{param_name}={default}"


def sig_params_to_str(sig, param_names, api_globals, indent=0):
    pos_only = []
    pos_or_kw = []
    var_positional = None
    kw_only = []
    var_keyword = None
    for param_name in param_names:
        param = sig.parameters[param_name]
        if param.kind == inspect.Parameter.POSITIONAL_ONLY:
            pos_only.append((param_name, param))
        elif param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
            pos_or_kw.append((param_name, param))
        elif param.kind == inspect.Parameter.VAR_POSITIONAL:
            var_positional = param_name
        elif param.kind == inspect.Parameter.KEYWORD_ONLY:
            kw_only.append((param_name, param))
        elif param.kind == inspect.Parameter.VAR_KEYWORD:
            var_keyword = param_name

    func_params = []
    body_params = []
    for param_name, param in pos_only:
        body_params.append(param_name)
        func_params.append(_add_typehint(param_name, param, api_globals))

    if pos_only:
        func_params.append("/")

    for param_name, param in pos_or_kw:
        body_par = param_name
        if not var_positional:
            body_par = f"{param_name}={param_name}"
        body_params.append(body_par)
        func_params.append(_kw_default_to_str(param_name, param, api_globals))

    if var_positional:
        body_params.append(f"*{var_positional}")
        func_params.append(f"*{var_positional}")

    elif kw_only:
        func_params.append("*")

    for param_name, param in kw_only:
        body_params.append(f"{param_name}={param_name}")
        func_params.append(_kw_default_to_str(param_name, param, api_globals))

    if var_keyword is not None:
        body_params.append(f"**{var_keyword}")
        func_params.append(f"**{var_keyword}")

    base_indent_str = " " * indent
    param_indent_str = " " * (indent + 4)

    func_params_str = "()"
    if func_params:
        lines_str = "\n".join([
            f"{param_indent_str}{line},"
            for line in func_params
        ])
        func_params_str = f"(\n{lines_str}\n{base_indent_str})"

    if sig.return_annotation is not inspect.Signature.empty:
        return_typehint = _get_typehint(sig.return_annotation, api_globals)
        func_params_str += f" -> {return_typehint}"

    body_params_str = "()"
    if body_params:
        lines_str = "\n".join([
            f"{param_indent_str}{line},"
            for line in body_params
        ])
        body_params_str = f"(\n{lines_str}\n{base_indent_str})"

    return func_params_str, body_params_str


def prepare_api_functions(api_globals):
    from ayon_api.server_api import (  # noqa: E402
        ServerAPI,
        InstallersAPI,
        DependencyPackagesAPI,
        SecretsAPI,
        BundlesAddonsAPI,
        EventsAPI,
        AttributesAPI,
        ProjectsAPI,
        FoldersAPI,
        TasksAPI,
        ProductsAPI,
        VersionsAPI,
        RepresentationsAPI,
        WorkfilesAPI,
        ThumbnailsAPI,
        ActivitiesAPI,
        ActionsAPI,
        LinksAPI,
        ListsAPI,
    )

    functions = []
    _items = list(ServerAPI.__dict__.items())
    _items.extend(InstallersAPI.__dict__.items())
    _items.extend(DependencyPackagesAPI.__dict__.items())
    _items.extend(SecretsAPI.__dict__.items())
    _items.extend(ActionsAPI.__dict__.items())
    _items.extend(ActivitiesAPI.__dict__.items())
    _items.extend(BundlesAddonsAPI.__dict__.items())
    _items.extend(EventsAPI.__dict__.items())
    _items.extend(AttributesAPI.__dict__.items())
    _items.extend(ProjectsAPI.__dict__.items())
    _items.extend(FoldersAPI.__dict__.items())
    _items.extend(TasksAPI.__dict__.items())
    _items.extend(ProductsAPI.__dict__.items())
    _items.extend(VersionsAPI.__dict__.items())
    _items.extend(RepresentationsAPI.__dict__.items())
    _items.extend(WorkfilesAPI.__dict__.items())
    _items.extend(LinksAPI.__dict__.items())
    _items.extend(ListsAPI.__dict__.items())
    _items.extend(ThumbnailsAPI.__dict__.items())

    processed = set()
    for attr_name, attr in _items:
        if attr_name in processed:
            continue
        processed.add(attr_name)
        if (
            attr_name.startswith("_")
            or attr_name in EXCLUDED_METHODS
            or not callable(attr)
        ):
            continue

        sig = inspect.signature(attr)
        param_names = list(sig.parameters)
        if inspect.isfunction(attr):
            param_names.pop(0)

        func_def_params, func_body_params = sig_params_to_str(
            sig, param_names, api_globals
        )

        func_def = f"def {attr_name}{func_def_params}:\n"

        func_body_parts = []
        docstring = prepare_docstring(attr)
        if docstring:
            func_body_parts.append(docstring)

        func_body_parts.extend([
            "con = get_server_api_connection()",
            f"return con.{attr_name}{func_body_params}",
        ])

        func_body = indent_lines("\n".join(func_body_parts))
        full_def = func_def + func_body
        functions.append(full_def)
    return "\n\n\n".join(functions)


def main():
    print("Creating public API functions based on ServerAPI methods")
    # TODO order methods in some order
    ayon_api_root = os.path.join(CURRENT_DIR, "ayon_api")
    init_filepath = os.path.join(ayon_api_root, "__init__.py")
    api_filepath = os.path.join(ayon_api_root, "_api.py")

    print("(1/5) Reading current content of '_api.py'")
    with open(api_filepath, "r") as stream:
        old_content = stream.read()

    parts = old_content.split(AUTOMATED_COMMENT)
    if len(parts) == 1:
        raise RuntimeError(
            "Automated comment not found in '_api.py'"
        )
    if len(parts) > 2:
        raise RuntimeError(
            "Automated comment found multiple times in '_api.py'"
        )

    print("(2/5) Parsing current '__init__.py' content")
    formatting_init_content = prepare_init_without_api(init_filepath)

    # Read content of first part of `_api.py` to get global variables
    # - first with disabled type checking so other files from ayon_api are
    #   loded without any issues
    typing.TYPE_CHECKING = False
    api_globals = {"__name__": "ayon_api._api"}
    exec(parts[0], api_globals)

    # - second with enabled type checking to get all available types in the
    #   file
    # NOTE The file contains 'from __future__ import annotations' so any
    #   typehints can be used, but we should validate if are available.
    typing.TYPE_CHECKING = True
    api_globals = {"__name__": "ayon_api._api"}
    exec(parts[0], api_globals)

    for attr_name in dir(__builtins__):
        api_globals[attr_name] = getattr(__builtins__, attr_name)

    print("(3/5) Preparing functions body based on 'ServerAPI' class")
    result = prepare_api_functions(api_globals)

    print("(4/5) Store new functions body to '_api.py'")
    new_content = f"{parts[0]}{AUTOMATED_COMMENT}\n{result}"
    with open(api_filepath, "w") as stream:
        print(new_content, file=stream)

    # find all functions and classes available in '_api.py'
    func_regex = re.compile(r"^(def|class) (?P<name>[^\(]*)(\(|:).*")
    func_names = []
    for line in new_content.split("\n"):
        result = func_regex.search(line)
        if result:
            name = result.group("name")
            if name.startswith("_") or name in EXCLUDED_IMPORT_NAMES:
                continue
            func_names.append(name)

    print("(5/5) Updating imports in '__init__.py'")
    import_lines = ["from ._api import ("]
    for name in func_names:
        import_lines.append(f"    {name},")
    import_lines.append(")")

    all_lines = [
        f'    "{name}",'
        for name in func_names
    ]
    new_init_content = formatting_init_content.format(
        api_imports="\n".join(import_lines),
        all_content="\n".join(all_lines),
    )

    with open(init_filepath, "w") as stream:
        print(new_init_content, file=stream)

    print("Public API functions created successfully")


if __name__ == "__main__":
    main()