from __future__ import annotations

import keyword

import re

from pathlib import Path

from jinja2 import Environment, FileSystemLoader, StrictUndefined

from .extensions import GeneratorExtensions

from .model import (

AnyAnnotation,

DictAnnotation,

EnumDef,

FieldDef,

GeneratedArtifact,

ListAnnotation,

LiteralAnnotation,

MappingAnnotation,

NamedAnnotation,

NormalizedSpec,

OperationDef,

TupleAnnotation,

TypeAliasDef,

TypeAnnotation,

TypedDictDef,

UnionAnnotation,

)

def _render_annotation(annotation: TypeAnnotation) -> str:

match annotation:

case AnyAnnotation():

return "Any"

case DictAnnotation(key, value):

return f"dict[{_render_annotation(key)}, {_render_annotation(value)}]"

case ListAnnotation(item):

return f"list[{_render_annotation(item)}]"

case MappingAnnotation(key, value):

return f"Mapping[{_render_annotation(key)}, {_render_annotation(value)}]"

case LiteralAnnotation(values):

return f"Literal[{', '.join(repr(value) for value in values)}]"

case NamedAnnotation(name):

return name

case TupleAnnotation(items):

return f"tuple[{', '.join(_render_annotation(item) for item in items)}]"

case UnionAnnotation(items):

return " | ".join(_render_annotation(item) for item in items) or "Any"

def _field_annotation(field: FieldDef) -> str:

"""

Jinja2 filter to format a FieldDef's annotation.

"""

annotation = repr(_render_field_annotation(field))

if not field.required:

annotation = f"NotRequired[{annotation}]"

return annotation

def _class_field_annotation(field: FieldDef, total_optional: bool) -> str:

annotation = _render_field_annotation(field)

if not field.required and not total_optional:

annotation = f"NotRequired[{annotation}]"

return annotation

def _render_field_annotation(field: FieldDef) -> str:

annotation = _render_annotation(field.annotation)

if field.description is None:

return annotation

return f"Annotated[{annotation}, _openapi_python_field({field.description!r})]"

def _string_literal(value: str) -> str:

return repr(value)

def _comment(value: str, spaces: int = 0) -> str:

prefix = " " * spaces

return "\n".join(

f"{prefix}# {line}" if line else f"{prefix}#" for line in value.splitlines()

)

def _supports_typeddict_class_syntax(defn: TypedDictDef) -> bool:

return all(

field.name.isidentifier()

and not keyword.iskeyword(field.name)

and not field.name.startswith("__")

for field in defn.fields

)

_TEMPLATE_DIR = Path(__file__).with_name("templates")

_JINJA_ENV = Environment(

loader=FileSystemLoader(_TEMPLATE_DIR),

trim_blocks=True,

lstrip_blocks=True,

undefined=StrictUndefined,

)

_JINJA_ENV.filters["repr"] = repr

_JINJA_ENV.filters["annotation"] = _render_annotation

_JINJA_ENV.filters["field_annotation"] = _field_annotation

_JINJA_ENV.filters["class_field_annotation"] = _class_field_annotation

_JINJA_ENV.filters["comment"] = _comment

_JINJA_ENV.filters["string_literal"] = _string_literal

def _render_template(name: str, **context: object) -> str:

return _JINJA_ENV.get_template(name).render(**context)

def _indent(text: str, spaces: int = 4) -> str:

prefix = " " * spaces

return "\n".join((prefix + line) if line else "" for line in text.splitlines())

def _format_typeddict(defn: TypedDictDef) -> str:

total_optional = bool(defn.fields) and all(

not field.required for field in defn.fields

)

return _render_template(

"typeddict.py.j2",

defn=defn,

class_syntax=_supports_typeddict_class_syntax(defn),

total_optional=total_optional,

)

def _format_alias(alias: TypeAliasDef) -> str:

return _render_template("alias.py.j2", alias=alias)

def _enum_member_name(value: object) -> str:

text = str(value).upper()

text = re.sub(r"[^a-zA-Z0-9_]", "_", text).strip("_") or "VALUE"

text = re.sub(r"_+", "_", text)

if text[0].isdigit():

text = f"_{text}"

if keyword.iskeyword(text.lower()):

text = f"{text}_"

return text

def _enum_base(defn: EnumDef) -> str:

if all(isinstance(value, str) for value in defn.values):

return "str, Enum"

if all(

isinstance(value, int) and not isinstance(value, bool) for value in defn.values

):

return "int, Enum"

return "Enum"

def _format_enum(defn: EnumDef) -> str:

members: list[tuple[str, object]] = []

used: set[str] = set()

for value in defn.values:

name = _enum_member_name(value)

base_name = name

index = 2

while name in used:

name = f"{base_name}_{index}"

index += 1

used.add(name)

members.append((name, value))

return _render_template(

"enum.py.j2",

defn=defn,

enum_base=_enum_base(defn),

members=members,

)

def _annotation_dependencies(annotation: TypeAnnotation, names: set[str]) -> set[str]:

match annotation:

case AnyAnnotation() | LiteralAnnotation():

return set()

case DictAnnotation(key, value):

return _annotation_dependencies(key, names) | _annotation_dependencies(

value, names

)

case MappingAnnotation(key, value):

return _annotation_dependencies(key, names) | _annotation_dependencies(

value, names

)

case ListAnnotation(item):

return _annotation_dependencies(item, names)

case NamedAnnotation(name):

return {name} if name in names else set()

case TupleAnnotation(items) | UnionAnnotation(items):

dependencies: set[str] = set()

for item in items:

dependencies.update(_annotation_dependencies(item, names))

return dependencies

def _typed_dict_dependencies(defn: TypedDictDef, names: set[str]) -> set[str]:

dependencies: set[str] = set()

for field in defn.fields:

dependencies.update(_annotation_dependencies(field.annotation, names))

dependencies.discard(defn.name)

return dependencies

def _alias_dependencies(defn: TypeAliasDef, names: set[str]) -> set[str]:

dependencies = _annotation_dependencies(defn.annotation, names)

dependencies.discard(defn.name)

return dependencies

def _type_dependencies(defn: TypeAliasDef | TypedDictDef, names: set[str]) -> set[str]:

match defn:

case TypeAliasDef():

return _alias_dependencies(defn, names)

case TypedDictDef():

return _typed_dict_dependencies(defn, names)

def _order_type_definitions(

aliases: tuple[TypeAliasDef, ...], typed_dicts: tuple[TypedDictDef, ...]

) -> list[TypeAliasDef | TypedDictDef]:

by_name: dict[str, TypeAliasDef | TypedDictDef] = {

item.name: item for item in (*aliases, *typed_dicts)

}

names = set(by_name)

ordered: list[TypeAliasDef | TypedDictDef] = []

temporary: set[str] = set()

permanent: set[str] = set()

def visit(name: str) -> None:

if name in permanent or name in temporary:

return

temporary.add(name)

for dependency in sorted(_type_dependencies(by_name[name], names)):

visit(dependency)

temporary.remove(name)

permanent.add(name)

ordered.append(by_name[name])

for name in sorted(by_name):

visit(name)

return ordered

def _format_type_definition(defn: TypeAliasDef | TypedDictDef) -> str:

match defn:

case TypeAliasDef():

return _format_alias(defn)

case TypedDictDef():

return _format_typeddict(defn)

def _has_field_descriptions(defns: tuple[TypedDictDef, ...]) -> bool:

return any(field.description is not None for defn in defns for field in defn.fields)

def _call_parameters(op: OperationDef, *, generate_requests: bool) -> dict[str, str]:

if not generate_requests:

return {

"params": "params: dict[str, Any] | None = None",

"query": "query: dict[str, Any] | None = None",

"headers": "headers: dict[str, Any] | None = None",

"body": "body: object | None = None",

}

params = "params: " + _render_annotation(op.params_type)

if not op.params_required:

params += " | None = None"

query = "query: " + _render_annotation(op.query_type)

if not op.query_required:

query += " | None = None"

headers = "headers: " + _render_annotation(op.headers_type)

if not op.headers_required:

headers += " | None = None"

body = "body: object | None = None"

if op.body_type is not None:

body = f"body: {_render_annotation(op.body_type)}"

if not op.body_required:

body += " | None = None"

return {

"params": params,

"query": query,

"headers": headers,

"body": body,

}

def _protocol_name(op: OperationDef, *, is_async: bool = False) -> str:

return f"Async{op.protocol_name}" if is_async else op.protocol_name

def _protocol_block(

op: OperationDef,

*,

generate_requests: bool,

generate_responses: bool,

is_async: bool = False,

) -> str:

return _render_template(

"protocol.py.j2",

op=op,

is_async=is_async,

protocol_name=_protocol_name(op, is_async=is_async),

call_parameters=_call_parameters(op, generate_requests=generate_requests),

response_type=(

_render_annotation(op.response_type) if generate_responses else "Any"

),

)

def _method_overload_line(

op: OperationDef, *, return_type: str, is_async: bool = False

) -> str:

return _render_template(

"method_overload.py.j2",

op=op,

return_type=return_type,

protocol_name=_protocol_name(op, is_async=is_async),

)

def _fallback_method_block(

method: str, overloads: list[str], *, is_async: bool = False

) -> str:

return _render_template(

"method_block.py.j2",

method=method,

overloads="\n".join(overloads),

callable_return="Awaitable[Any]" if is_async else "object",

call_return="Any" if is_async else "object",

is_async=is_async,

)

def _included_annotations(

spec: NormalizedSpec, *, generate_requests: bool, generate_responses: bool

) -> tuple[TypeAnnotation, ...]:

"""

Returns the set of type annotations that be included

in the rendered output.

"""

roots: list[TypeAnnotation] = []

for op in spec.operations:

if generate_requests:

roots.extend((op.params_type, op.query_type, op.headers_type))

if op.body_type is not None:

roots.append(op.body_type)

if generate_responses:

roots.append(op.response_type)

return tuple(roots)

def _used_type_names(

spec: NormalizedSpec, *, generate_requests: bool, generate_responses: bool

) -> set[str]:

"""

Returns the set of type definition names that are transitively referenced by

the client protocols.

"""

by_name: dict[str, TypeAliasDef | TypedDictDef | EnumDef] = {

item.name: item for item in (*spec.aliases, *spec.typed_dicts, *spec.enums)

}

all_names = set(by_name)

used: set[str] = set()

pending: list[str] = []

for annotation in _included_annotations(

spec,

generate_requests=generate_requests,

generate_responses=generate_responses,

):

pending.extend(_annotation_dependencies(annotation, all_names) - used)

while pending:

name = pending.pop()

if name in used:

continue

used.add(name)

item = by_name[name]

match item:

case TypeAliasDef() | TypedDictDef():

pending.extend(_type_dependencies(item, all_names) - used)

case EnumDef():

pass

return used

def _render_types(

spec: NormalizedSpec,

*,

generate_routes: bool,

generate_requests: bool,

generate_responses: bool,

) -> str:

route_aliases = _route_aliases(spec, generate_routes=generate_routes)

used_names = _used_type_names(

spec,

generate_requests=generate_requests,

generate_responses=generate_responses,

)

aliases = tuple(item for item in spec.aliases if item.name in used_names)

typed_dicts = tuple(item for item in spec.typed_dicts if item.name in used_names)

enums = tuple(item for item in spec.enums if item.name in used_names)

type_definitions = _order_type_definitions((*route_aliases, *aliases), typed_dicts)

blocks = [_format_enum(item) for item in enums] + [

_format_type_definition(item) for item in type_definitions

]

return _render_template(

"types.py.j2",

has_field_descriptions=_has_field_descriptions(typed_dicts),

type_blocks="\n".join(blocks).strip() + "\n",

)

def rendered_type_definition_count(

spec: NormalizedSpec,

*,

generate_routes: bool,

generate_requests: bool,

generate_responses: bool,

) -> int:

used_names = _used_type_names(

spec,

generate_requests=generate_requests,

generate_responses=generate_responses,

)

return (

len(_route_aliases(spec, generate_routes=generate_routes))

+ sum(1 for item in spec.aliases if item.name in used_names)

+ sum(1 for item in spec.typed_dicts if item.name in used_names)

+ sum(1 for item in spec.enums if item.name in used_names)

)

def _literal_annotation(values: set[str]) -> LiteralAnnotation:

return LiteralAnnotation(tuple(sorted(values)))

def _route_aliases(

spec: NormalizedSpec, *, generate_routes: bool

) -> tuple[TypeAliasDef, ...]:

if not generate_routes:

return (TypeAliasDef(name="RouteLiteral", annotation=NamedAnnotation("str")),)

routes_by_method: dict[str, set[str]] = {}

for op in spec.operations:

routes_by_method.setdefault(op.method.upper(), set()).add(op.route_literal)

aliases = [

TypeAliasDef(

name=f"{method}_RouteLiteral", annotation=_literal_annotation(routes)

)

for method, routes in sorted(routes_by_method.items())

]

if aliases:

aliases.append(

TypeAliasDef(

name="RouteLiteral",

annotation=UnionAnnotation(

tuple(NamedAnnotation(alias.name) for alias in aliases)

),

)

)

else:

aliases.append(

TypeAliasDef(name="RouteLiteral", annotation=NamedAnnotation("str"))

)

return tuple(aliases)

def _render_transport(spec: NormalizedSpec, *, protocol_only: bool) -> str:

return _render_template(

"transport.py.j2",

typing_imports=("Protocol" if protocol_only else "TYPE_CHECKING, Protocol"),

include_default_transport=not protocol_only,

)

def _render_client(

spec: NormalizedSpec,

*,

protocol_only: bool,

generate_routes: bool,

generate_requests: bool,

generate_responses: bool,

) -> str:

protocols: list[str] = []

async_protocols: list[str] = []

method_overloads: dict[str, list[str]] = {}

async_method_overloads: dict[str, list[str]] = {}

generate_operation_protocols = generate_routes and (

generate_requests or generate_responses

)

for op in spec.operations:

if generate_operation_protocols:

protocols.append(

_protocol_block(

op,

generate_requests=generate_requests,

generate_responses=generate_responses,

)

)

async_protocols.append(

_protocol_block(

op,

generate_requests=generate_requests,

generate_responses=generate_responses,

is_async=True,

)

)

method_overloads.setdefault(op.method, []).append(

_method_overload_line(

op, return_type=_protocol_name(op), is_async=False

)

)

async_method_overloads.setdefault(op.method, []).append(

_method_overload_line(

op,

return_type=_protocol_name(op, is_async=True),

is_async=True,

)

)

else:

overloads = method_overloads.setdefault(op.method, [])

async_overloads = async_method_overloads.setdefault(op.method, [])

if generate_routes:

overloads.append(

_method_overload_line(

op, return_type="Callable[..., object]", is_async=False

)

)

async_overloads.append(

_method_overload_line(

op, return_type="Callable[..., Awaitable[Any]]", is_async=True

)

)

method_blocks: list[str] = []

for method in sorted(method_overloads):

method_blocks.append(

_fallback_method_block(

method,

method_overloads[method],

)

)

async_method_blocks: list[str] = []

for method in sorted(async_method_overloads):

async_method_blocks.append(

_fallback_method_block(

method,

async_method_overloads[method],

is_async=True,

)

)

if not protocol_only:

transport_imports = (

"from .transport import AsyncTransport, DefaultAsyncTransport, "

"DefaultTransport, Transport"

)

sync_transport_type = "Transport | None = None"

sync_transport_assignment = "transport or DefaultTransport()"

async_transport_type = "AsyncTransport | None = None"

async_transport_assignment = "transport or DefaultAsyncTransport()"

else:

transport_imports = "from .transport import AsyncTransport, Transport"

sync_transport_type = "Transport"

sync_transport_assignment = "transport"

async_transport_type = "AsyncTransport"

async_transport_assignment = "transport"

return _render_template(

"client.py.j2",

transport_imports=transport_imports,

sync_transport_type=sync_transport_type,

sync_transport_assignment=sync_transport_assignment,

async_transport_type=async_transport_type,

async_transport_assignment=async_transport_assignment,

protocol_blocks="\n".join(protocols).strip() + "\n",

async_protocol_blocks="\n".join(async_protocols).strip() + "\n",

method_blocks=_indent("\n".join(method_blocks).strip() + "\n"),

async_method_blocks=_indent("\n".join(async_method_blocks).strip() + "\n"),

)

def render_package(

spec: NormalizedSpec,

extensions: GeneratorExtensions | None = None,

*,

protocol_only: bool = False,

generate_routes: bool = True,

generate_requests: bool = True,

generate_responses: bool = True,

) -> list[GeneratedArtifact]:

context = {

"types": _render_types(

spec,

generate_routes=generate_routes,

generate_requests=generate_requests,

generate_responses=generate_responses,

),

"transport": _render_transport(spec, protocol_only=protocol_only),

"client": _render_client(

spec,

protocol_only=protocol_only,

generate_routes=generate_routes,

generate_requests=generate_requests,

generate_responses=generate_responses,

),

}

if extensions:

for hook in extensions.render_context_hooks:

context = hook(spec, context)

init_content = _render_template(

"init.py.j2", include_default_transport=not protocol_only

)

return [

GeneratedArtifact(

relative_path=f"{spec.package_name}/__init__.py", content=init_content

),

GeneratedArtifact(

relative_path=f"{spec.package_name}/types.py", content=context["types"]

),

GeneratedArtifact(

relative_path=f"{spec.package_name}/transport.py",

content=context["transport"],

),

GeneratedArtifact(

relative_path=f"{spec.package_name}/client.py", content=context["client"]

),

GeneratedArtifact(relative_path=f"{spec.package_name}/py.typed", content="\n"),

]