Skip to content

Commit

Permalink
Limiting FILE_HEADER (#114)
Browse files Browse the repository at this point in the history
Why
===

Now that we're generating modules, not all imports apply in all places,
and it imports are not able to be pruned in `__init__.py`, so let's not
just import a bunch of stuff we don't need.

What changed
============

Breaking `FILE_HEADER` up into three different variants.

Test plan
=========

_Describe what you did to test this change to a level of detail that
allows your reviewer to test it_
  • Loading branch information
blast-hardcheese authored Nov 21, 2024
1 parent d85becf commit fbe67b7
Showing 1 changed file with 28 additions and 2 deletions.
30 changes: 28 additions & 2 deletions replit_river/codegen/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,32 @@ def ensure_literal_type(value: TypeExpression) -> TypeName:

_NON_ALNUM_RE = re.compile(r"[^a-zA-Z0-9_]+")

# Literal is here because HandshakeType can be Literal[None]
ROOT_FILE_HEADER = dedent(
"""\
# Code generated by river.codegen. DO NOT EDIT.
from pydantic import BaseModel
from typing import Literal
import replit_river as river
"""
)

SERVICE_FILE_HEADER = dedent(
"""\
# Code generated by river.codegen. DO NOT EDIT.
from collections.abc import AsyncIterable, AsyncIterator
from typing import Any
from pydantic import TypeAdapter
from replit_river.error_schema import RiverError
import replit_river as river
"""
)

FILE_HEADER = dedent(
"""\
# ruff: noqa
Expand Down Expand Up @@ -709,7 +735,7 @@ def generate_common_client(
handshake_chunks: Sequence[str],
modules: list[Tuple[ModuleName, ClassName]],
) -> FileContents:
chunks: list[str] = [FILE_HEADER]
chunks: list[str] = [ROOT_FILE_HEADER]
chunks.extend(
[
f"from .{model_name} import {class_name}"
Expand Down Expand Up @@ -1072,7 +1098,7 @@ async def {name}(
]

emitted_files[RenderedPath(str(Path(f"{schema_name}/__init__.py")))] = FileContents(
"\n".join([FILE_HEADER] + rendered_imports + in_root + current_chunks)
"\n".join([SERVICE_FILE_HEADER] + rendered_imports + in_root + current_chunks)
)
return (
ModuleName(schema_name),
Expand Down

0 comments on commit fbe67b7

Please sign in to comment.