code: #  Copyright 2023 LanceDB Developers
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.
  
  
  import json
  import sys
  from datetime import date, datetime
  from typing import List, Optional, Tuple
  
  import pyarrow as pa
  import pydantic
  import pytest
  from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema
  from pydantic import Field
  
  
  @pytest.mark.skipif(
      sys.version_info < (3, 9),
      reason="using native type alias requires python3.9 or higher",
  )
  def test_pydantic_to_arrow():
      class StructModel(pydantic.BaseModel):
          a: str
          b: Optional[float]
  
      class TestModel(pydantic.BaseModel):
          id: int
          s: str
          vec: list[float]
          li: list[int]
          lili: list[list[float]]
          litu: list[tuple[float, float]]
          opt: Optional[str] = None
          st: StructModel
          dt: date
          dtt: datetime
          dt_with_tz: datetime = Field(json_schema_extra={"tz": "Asia/Shanghai"})
          # d: dict
  
      # TODO: test we can actually convert the model into data.
      # m = TestModel(
      #     id=1,
      #     s="hello",
      #     vec=[1.0, 2.0, 3.0],
      #     li=[2, 3, 4],
      #     lili=[[2.5, 1.5], [3.5, 4.5], [5.5, 6.5]],
      #     litu=[(2.5, 1.5), (3.5, 4.5), (5.5, 6.5)],
      #     st=StructModel(a="a", b=1.0),
      #     dt=date.today(),
      #     dtt=datetime.now(),
      #     dt_with_tz=datetime.now(pytz.timezone("Asia/Shanghai")),
      # )
  
      schema = pydantic_to_schema(TestModel)
  
      expect_schema = pa.schema(
          [
              pa.field("id", pa.int64(), False),
              pa.field("s", pa.utf8(), False),
              pa.field("vec", pa.list_(pa.float64()), False),
              pa.field("li", pa.list_(pa.int64()), False),
              pa.field("lili", pa.list_(pa.list_(pa.float64())), False),
              pa.field("litu", pa.list_(pa.list_(pa.float64())), False),
              pa.field("opt", pa.utf8(), True),
              pa.field(
                  "st",
                  pa.struct(
                      [pa.field("a", pa.utf8(), False), pa.field("b", pa.float64(), True)]
                  ),
                  False,
              ),
              pa.field("dt", pa.date32(), False),
              pa.field("dtt", pa.timestamp("us"), False),
              pa.field("dt_with_tz", pa.timestamp("us", tz="Asia/Shanghai"), False),
          ]
      )
      assert schema == expect_schema
  
  
  @pytest.mark.skipif(
      sys.version_info < (3, 10),
      reason="using | type syntax requires python3.10 or higher",
  )
  def test_optional_types_py310():
      class TestModel(pydantic.BaseModel):
          a: str | None
          b: None | str
          c: Optional[str]
  
      schema = pydantic_to_schema(TestModel)
  
      expect_schema = pa.schema(
          [
              pa.field("a", pa.utf8(), True),
              pa.field("b", pa.utf8(), True),
              pa.field("c", pa.utf8(), True),
          ]
      )
      assert schema == expect_schema
  
  
  @pytest.mark.skipif(
      sys.version_info > (3, 8),
      reason="using native type alias requires python3.9 or higher",
  )
  def test_pydantic_to_arrow_py38():
      class StructModel(pydantic.BaseModel):
          a: str
          b: Optional[float]
  
      class TestModel(pydantic.BaseModel):
          id: int
          s: str
          vec: List[float]
          li: List[int]
          lili: List[List[float]]
          litu: List[Tuple[float, float]]
          opt: Optional[str] = None
          st: StructModel
          dt: date
          dtt: datetime
          dt_with_tz: datetime = Field(json_schema_extra={"tz": "Asia/Shanghai"})
          # d: dict
  
      # TODO: test we can actually convert the model to Arrow data.
      # m = TestModel(
      #     id=1,
      #     s="hello",
      #     vec=[1.0, 2.0, 3.0],
      #     li=[2, 3, 4],
      #     lili=[[2.5, 1.5], [3.5, 4.5], [5.5, 6.5]],
      #     litu=[(2.5, 1.5), (3.5, 4.5), (5.5, 6.5)],
      #     st=StructModel(a="a", b=1.0),
      #     dt=date.today(),
      #     dtt=datetime.now(),
      #     dt_with_tz=datetime.now(pytz.timezone("Asia/Shanghai")),
      # )
  
      schema = pydantic_to_schema(TestModel)
  
      expect_schema = pa.schema(
          [
              pa.field("id", pa.int64(), False),
              pa.field("s", pa.utf8(), False),
              pa.field("vec", pa.list_(pa.float64()), False),
              pa.field("li", pa.list_(pa.int64()), False),
              pa.field("lili", pa.list_(pa.list_(pa.float64())), False),
              pa.field("litu", pa.list_(pa.list_(pa.float64())), False),
              pa.field("opt", pa.utf8(), True),
              pa.field(
                  "st",
                  pa.struct(
                      [pa.field("a", pa.utf8(), False), pa.field("b", pa.float64(), True)]
                  ),
                  False,
              ),
              pa.field("dt", pa.date32(), False),
              pa.field("dtt", pa.timestamp("us"), False),
              pa.field("dt_with_tz", pa.timestamp("us", tz="Asia/Shanghai"), False),
          ]
      )
      assert schema == expect_schema
  
  
  def test_fixed_size_list_field():
      class TestModel(pydantic.BaseModel):
          vec: Vector(16)
          li: List[int]
  
      data = TestModel(vec=list(range(16)), li=[1, 2, 3])
      if PYDANTIC_VERSION >= (2,):
          assert json.loads(data.model_dump_json()) == {
              "vec": list(range(16)),
              "li": [1, 2, 3],
          }
      else:
          assert data.dict() == {
              "vec": list(range(16)),
              "li": [1, 2, 3],
          }
  
      schema = pydantic_to_schema(TestModel)
      assert schema == pa.schema(
          [
              pa.field("vec", pa.list_(pa.float32(), 16), False),
              pa.field("li", pa.list_(pa.int64()), False),
          ]
      )
  
      if PYDANTIC_VERSION >= (2,):
          json_schema = TestModel.model_json_schema()
      else:
          json_schema = TestModel.schema()
  
      assert json_schema == {
          "properties": {
              "vec": {
                  "items": {"type": "number"},
                  "maxItems": 16,
                  "minItems": 16,
                  "title": "Vec",
                  "type": "array",
              },
              "li": {"items": {"type": "integer"}, "title": "Li", "type": "array"},
          },
          "required": ["vec", "li"],
          "title": "TestModel",
          "type": "object",
      }
  
  
  def test_fixed_size_list_validation():
      class TestModel(pydantic.BaseModel):
          vec: Vector(8)
  
      with pytest.raises(pydantic.ValidationError):
          TestModel(vec=range(9))
  
      with pytest.raises(pydantic.ValidationError):
          TestModel(vec=range(7))
  
      TestModel(vec=range(8))
  
  
  def test_lance_model():
      class TestModel(LanceModel):
          vector: Vector(16) = Field(default=[0.0] * 16)
          li: List[int] = Field(default=[1, 2, 3])
  
      schema = pydantic_to_schema(TestModel)
      assert schema == TestModel.to_arrow_schema()
      assert TestModel.field_names() == ["vector", "li"]
  
      t = TestModel()
      assert t == TestModel(vec=[0.0] * 16, li=[1, 2, 3])
  
apis: lancedb.pydantic.Vector lancedb.pydantic.pydantic_to_schema
extract_api: [((860, 973), 'pytest.mark.skipif', 'pytest.mark.skipif', (['(sys.version_info < (3, 9))'], {'reason': '"""using native type alias requires python3.9 or higher"""'}), "(sys.version_info < (3, 9), reason=\n    'using native type alias requires python3.9 or higher')\n", (878, 973), False, 'import pytest\n'), ((2877, 2988), 'pytest.mark.skipif', 'pytest.mark.skipif', (['(sys.version_info < (3, 10))'], {'reason': '"""using | type syntax requires python3.10 or higher"""'}), "(sys.version_info < (3, 10), reason=\n    'using | type syntax requires python3.10 or higher')\n", (2895, 2988), False, 'import pytest\n'), ((3410, 3523), 'pytest.mark.skipif', 'pytest.mark.skipif', (['(sys.version_info > (3, 8))'], {'reason': '"""using native type alias requires python3.9 or higher"""'}), "(sys.version_info > (3, 8), reason=\n    'using native type alias requires python3.9 or higher')\n", (3428, 3523), False, 'import pytest\n'), ((1950, 1979), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['TestModel'], {}), '(TestModel)\n', (1968, 1979), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((3152, 3181), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['TestModel'], {}), '(TestModel)\n', (3170, 3181), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((4509, 4538), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['TestModel'], {}), '(TestModel)\n', (4527, 4538), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((5907, 5936), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['TestModel'], {}), '(TestModel)\n', (5925, 5936), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((7183, 7212), 'lancedb.pydantic.pydantic_to_schema', 'pydantic_to_schema', (['TestModel'], {}), '(TestModel)\n', (7201, 7212), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((1415, 1463), 'pydantic.Field', 'Field', ([], {'json_schema_extra': "{'tz': 'Asia/Shanghai'}"}), "(json_schema_extra={'tz': 'Asia/Shanghai'})\n", (1420, 1463), False, 'from pydantic import Field\n'), ((3970, 4018), 'pydantic.Field', 'Field', ([], {'json_schema_extra': "{'tz': 'Asia/Shanghai'}"}), "(json_schema_extra={'tz': 'Asia/Shanghai'})\n", (3975, 4018), False, 'from pydantic import Field\n'), ((5523, 5533), 'lancedb.pydantic.Vector', 'Vector', (['(16)'], {}), '(16)\n', (5529, 5533), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((6801, 6810), 'lancedb.pydantic.Vector', 'Vector', (['(8)'], {}), '(8)\n', (6807, 6810), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((6821, 6860), 'pytest.raises', 'pytest.raises', (['pydantic.ValidationError'], {}), '(pydantic.ValidationError)\n', (6834, 6860), False, 'import pytest\n'), ((6904, 6943), 'pytest.raises', 'pytest.raises', (['pydantic.ValidationError'], {}), '(pydantic.ValidationError)\n', (6917, 6943), False, 'import pytest\n'), ((7081, 7091), 'lancedb.pydantic.Vector', 'Vector', (['(16)'], {}), '(16)\n', (7087, 7091), False, 'from lancedb.pydantic import PYDANTIC_VERSION, LanceModel, Vector, pydantic_to_schema\n'), ((7094, 7119), 'pydantic.Field', 'Field', ([], {'default': '([0.0] * 16)'}), '(default=[0.0] * 16)\n', (7099, 7119), False, 'from pydantic import Field\n'), ((7144, 7168), 'pydantic.Field', 'Field', ([], {'default': '[1, 2, 3]'}), '(default=[1, 2, 3])\n', (7149, 7168), False, 'from pydantic import Field\n'), ((2049, 2059), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (2057, 2059), True, 'import pyarrow as pa\n'), ((2095, 2104), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (2102, 2104), True, 'import pyarrow as pa\n'), ((2401, 2410), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (2408, 2410), True, 'import pyarrow as pa\n'), ((2663, 2674), 'pyarrow.date32', 'pa.date32', ([], {}), '()\n', (2672, 2674), True, 'import pyarrow as pa\n'), ((2712, 2730), 'pyarrow.timestamp', 'pa.timestamp', (['"""us"""'], {}), "('us')\n", (2724, 2730), True, 'import pyarrow as pa\n'), ((2775, 2813), 'pyarrow.timestamp', 'pa.timestamp', (['"""us"""'], {'tz': '"""Asia/Shanghai"""'}), "('us', tz='Asia/Shanghai')\n", (2787, 2813), True, 'import pyarrow as pa\n'), ((3250, 3259), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (3257, 3259), True, 'import pyarrow as pa\n'), ((3294, 3303), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (3301, 3303), True, 'import pyarrow as pa\n'), ((3338, 3347), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (3345, 3347), True, 'import pyarrow as pa\n'), ((4608, 4618), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (4616, 4618), True, 'import pyarrow as pa\n'), ((4654, 4663), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (4661, 4663), True, 'import pyarrow as pa\n'), ((4960, 4969), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (4967, 4969), True, 'import pyarrow as pa\n'), ((5222, 5233), 'pyarrow.date32', 'pa.date32', ([], {}), '()\n', (5231, 5233), True, 'import pyarrow as pa\n'), ((5271, 5289), 'pyarrow.timestamp', 'pa.timestamp', (['"""us"""'], {}), "('us')\n", (5283, 5289), True, 'import pyarrow as pa\n'), ((5334, 5372), 'pyarrow.timestamp', 'pa.timestamp', (['"""us"""'], {'tz': '"""Asia/Shanghai"""'}), "('us', tz='Asia/Shanghai')\n", (5346, 5372), True, 'import pyarrow as pa\n'), ((2151, 2163), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2161, 2163), True, 'import pyarrow as pa\n'), ((2210, 2220), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (2218, 2220), True, 'import pyarrow as pa\n'), ((4710, 4722), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (4720, 4722), True, 'import pyarrow as pa\n'), ((4769, 4779), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (4777, 4779), True, 'import pyarrow as pa\n'), ((2278, 2290), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2288, 2290), True, 'import pyarrow as pa\n'), ((2349, 2361), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2359, 2361), True, 'import pyarrow as pa\n'), ((4837, 4849), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (4847, 4849), True, 'import pyarrow as pa\n'), ((4908, 4920), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (4918, 4920), True, 'import pyarrow as pa\n'), ((6016, 6028), 'pyarrow.float32', 'pa.float32', ([], {}), '()\n', (6026, 6028), True, 'import pyarrow as pa\n'), ((6079, 6089), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (6087, 6089), True, 'import pyarrow as pa\n'), ((2525, 2534), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (2532, 2534), True, 'import pyarrow as pa\n'), ((2558, 2570), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (2568, 2570), True, 'import pyarrow as pa\n'), ((5084, 5093), 'pyarrow.utf8', 'pa.utf8', ([], {}), '()\n', (5091, 5093), True, 'import pyarrow as pa\n'), ((5117, 5129), 'pyarrow.float64', 'pa.float64', ([], {}), '()\n', (5127, 5129), True, 'import pyarrow as pa\n')]