From 8f1f3741ef32df57921922368806fc604d273669 Mon Sep 17 00:00:00 2001
From: Carol Jung <cjung@codegen.com>
Date: Wed, 29 Jan 2025 23:14:02 -0800
Subject: [PATCH] Update span to use v2 pydantic

---
 src/codegen/sdk/codebase/span.py              | 75 ++++++++++---------
 .../test_api_doc_generation.py                |  2 +-
 2 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/src/codegen/sdk/codebase/span.py b/src/codegen/sdk/codebase/span.py
index 1f8b6644f..799482fa9 100644
--- a/src/codegen/sdk/codebase/span.py
+++ b/src/codegen/sdk/codebase/span.py
@@ -1,6 +1,7 @@
 from typing import Annotated, Any
 
-from pydantic import BaseModel, ConfigDict, PlainSerializer, PlainValidator, WithJsonSchema
+from pydantic import BaseModel, ConfigDict, PlainValidator, WithJsonSchema
+from pydantic.json_schema import JsonSchemaValue
 from pydantic_core.core_schema import ValidationInfo
 from tree_sitter import Point, Range
 
@@ -21,43 +22,31 @@ def validate_range(value: Any, info: ValidationInfo) -> Range:
     return value
 
 
-RangeAdapter = Annotated[
-    Range,
-    PlainValidator(validate_range),
-    PlainSerializer(
-        lambda range: {
-            "start_byte": range.start_byte,
-            "end_byte": range.end_byte,
+def range_json_schema() -> JsonSchemaValue:
+    return {
+        "type": "object",
+        "properties": {
+            "start_byte": {"type": "integer"},
+            "end_byte": {"type": "integer"},
             "start_point": {
-                "row": range.start_point.row,
-                "column": range.start_point.column,
+                "type": "object",
+                "properties": {
+                    "row": {"type": "integer"},
+                    "column": {"type": "integer"},
+                },
             },
             "end_point": {
-                "row": range.end_point.row,
-                "column": range.end_point.column,
-            },
-        }
-    ),
-    WithJsonSchema(
-        {
-            "type": "object",
-            "properties": {
-                "start_byte": {"type": "integer"},
-                "end_byte": {"type": "integer"},
-                "start_point": {
-                    "type": "object",
-                    "properties": {
-                        "row": {"type": "integer"},
-                        "column": {"type": "integer"},
-                    },
-                },
-                "end_point": {
-                    "type": "object",
-                    "properties": {"row": {"type": "integer"}, "column": {"type": "integer"}},
-                },
+                "type": "object",
+                "properties": {"row": {"type": "integer"}, "column": {"type": "integer"}},
             },
-        }
-    ),
+        },
+    }
+
+
+RangeAdapter = Annotated[
+    Range,
+    PlainValidator(validate_range),
+    WithJsonSchema(range_json_schema()),
 ]
 
 
@@ -65,6 +54,22 @@ def validate_range(value: Any, info: ValidationInfo) -> Range:
 class Span(BaseModel):
     """Range within the codebase"""
 
-    model_config = ConfigDict(frozen=True)
+    model_config = ConfigDict(
+        frozen=True,
+        json_encoders={
+            Range: lambda r: {
+                "start_byte": r.start_byte,
+                "end_byte": r.end_byte,
+                "start_point": {
+                    "row": r.start_point.row,
+                    "column": r.start_point.column,
+                },
+                "end_point": {
+                    "row": r.end_point.row,
+                    "column": r.end_point.column,
+                },
+            }
+        },
+    )
     range: RangeAdapter
     filepath: str
diff --git a/tests/unit/codegen/sdk/code_generation/test_api_doc_generation.py b/tests/unit/codegen/sdk/code_generation/test_api_doc_generation.py
index 665be5b5b..d6e72cf85 100644
--- a/tests/unit/codegen/sdk/code_generation/test_api_doc_generation.py
+++ b/tests/unit/codegen/sdk/code_generation/test_api_doc_generation.py
@@ -47,7 +47,7 @@ def test_api_doc_generation_sanity(codebase, language: ProgrammingLanguage) -> N
     # assert "InviteFactoryCreateParams" in docs # Canonicals aren't in docs
 
 
-@pytest.mark.timeout(120)
+@pytest.mark.timeout(160)
 @pytest.mark.xdist_group("codegen")
 def test_mdx_api_doc_generation_sanity(codebase) -> None:
     docs_json = generate_docs_json(codebase, "HEAD")