Limiting FILE_HEADER

Author: blast-hardcheese

replit_river/codegen/client.py

@@ -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
@@ -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}"
@@ -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),