def generate_future_compatible_method()

in utils/_legacy_check_future_compatible_signatures.py [0:0]

• 41 lines of code
• 9 McCabe index (conditional complexity)

def generate_future_compatible_method(method: Callable, method_source: str) -> str:
    # 1. Document `as_future` parameter
    if AS_FUTURE_DOCSTRING_TEMPLATE not in method_source:
        match = ARGS_DOCSTRING_REGEX.search(method_source)
        if match is None:
            raise ValueError(f"Could not find `Args` section in docstring of {method}.")
        args_docs = match.group(1).strip()
        method_source = method_source.replace(args_docs, args_docs + AS_FUTURE_DOCSTRING_TEMPLATE)

    # 2. Update signature
    # 2.a. Add `as_future` parameter
    if AS_FUTURE_SIGNATURE_TEMPLATE not in method_source:
        match = SIGNATURE_REGEX_FULL.search(method_source)
        if match is None:
            raise ValueError(f"Could not find signature of {method} in source.")
        method_source = method_source.replace(
            match.group(), match.group().replace(") ->", f" {AS_FUTURE_SIGNATURE_TEMPLATE}) ->"), 1
        )

    # 2.b. Update return value
    if "Future[" not in method_source:
        match = SIGNATURE_REGEX_RETURN_TYPE.search(method_source)
        if match is None:
            raise ValueError(f"Could not find return type of {method} in source.")
        base_type = match.group(1).strip()
        return_type = f"Union[{base_type}, Future[{base_type}]]"
        return_value_replaced = match.group().replace(match.group(1), return_type)
        method_source = method_source.replace(match.group(), return_value_replaced)

    # 3. Generate @overload stubs
    match = SIGNATURE_REGEX_FULL.search(method_source)
    if match is None:
        raise ValueError(f"Could not find signature of {method} in source.")
    method_sig = match.group()

    match = SIGNATURE_REGEX_RETURN_TYPE_WITH_FUTURE.search(method_sig)
    if match is None:
        raise ValueError(f"Could not find return type (with Future) of {method} in source.")
    no_future_return_type = match.group(1).strip()
    with_future_return_type = match.group(2).strip()

    # 3.a. Stub when `as_future=False`
    no_future_stub = "    @overload\n" + method_sig
    no_future_stub = no_future_stub.replace(AS_FUTURE_SIGNATURE_TEMPLATE, "as_future: Literal[False] = ...")
    no_future_stub = SIGNATURE_REGEX_RETURN_TYPE.sub(rf"-> {no_future_return_type}:", no_future_stub)
    no_future_stub += "  # type: ignore\n        ..."  # only the first stub requires "type: ignore"

    # 3.b. Stub when `as_future=True`
    with_future_stub = "    @overload\n" + method_sig
    with_future_stub = with_future_stub.replace(AS_FUTURE_SIGNATURE_TEMPLATE, "as_future: Literal[True] = ...")
    with_future_stub = SIGNATURE_REGEX_RETURN_TYPE.sub(rf"-> {with_future_return_type}:", with_future_stub)
    with_future_stub += "\n        ..."

    stubs_source = no_future_stub + "\n\n" + with_future_stub + "\n\n"

    # 4. All good!
    return method_source, stubs_source