working?

python/python/lancedb/pydantic.py

@@ -37,30 +37,26 @@ import numpy as np
 import pyarrow as pa
 import pydantic
 import semver
-from pydantic.fields import FieldInfo
 from lance.arrow import (
-    EncodedImageArray,
-    EncodedImageScalar,
     EncodedImageType,
-    ImageURIArray,
-    ImageURIScalar,
-    ImageURIType,
 )
+from pydantic.fields import FieldInfo
+from pydantic_core import core_schema
 
-PYDANTIC_VERSION = semver.Version.parse(pydantic.__version__)
-try:
-    from pydantic_core import CoreSchema, core_schema
-    from pydantic import GetJsonSchemaHandler
-    from pydantic.json_schema import JsonSchemaValue
-except ImportError:
-    if PYDANTIC_VERSION >= (2,):
-        raise
 
 if TYPE_CHECKING:
     from pydantic.fields import FieldInfo
-
     from .embeddings import EmbeddingFunctionConfig
 
+    PYDANTIC_VERSION = semver.Version.parse(pydantic.__version__)
+    try:
+        from pydantic import GetJsonSchemaHandler
+        from pydantic.json_schema import JsonSchemaValue
+        from pydantic_core import CoreSchema
+    except ImportError:
+        if PYDANTIC_VERSION >= (2,):
+            raise
+
 
 class FixedSizeListMixin(ABC):
     @staticmethod
@@ -135,7 +131,7 @@ def Vector(
         @classmethod
         def __get_pydantic_core_schema__(
             cls, _source_type: Any, _handler: pydantic.GetCoreSchemaHandler
-        ) -> CoreSchema:
+        ) -> "CoreSchema":
             return core_schema.no_info_after_validator_function(
                 cls,
                 core_schema.list_schema(
@@ -238,7 +234,7 @@ def EncodedImage():
         @classmethod
         def __get_pydantic_core_schema__(
             cls, _source_type: Any, _handler: pydantic.GetCoreSchemaHandler
-        ) -> CoreSchema:
+        ) -> "CoreSchema":
             from_bytes_schema = core_schema.bytes_schema()
 
             return core_schema.json_or_python_schema(
@@ -256,8 +252,8 @@ def EncodedImage():
 
         @classmethod
         def __get_pydantic_json_schema__(
-            cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+            cls, _core_schema: "CoreSchema", handler: "GetJsonSchemaHandler"
-        ) -> JsonSchemaValue:
+        ) -> "JsonSchemaValue":
             return handler(core_schema.bytes_schema())
 
         @classmethod

python/python/tests/test_cli.py

@@ -1,4 +1,5 @@
 from click.testing import CliRunner
+
 from lancedb.cli.cli import cli
 from lancedb.utils import CONFIG
 

python/python/tests/test_context.py

@@ -13,6 +13,7 @@
 
 import pandas as pd
 import pytest
+
 from lancedb.context import contextualize
 
 

python/python/tests/test_db.py

@@ -19,6 +19,8 @@ import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytest
+
+import lancedb
 from lancedb.pydantic import LanceModel, Vector
 
 

python/python/tests/test_e2e_remote_db.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import pytest
+
 from lancedb import LanceDBConnection
 
 # TODO: setup integ test mark and script

python/python/tests/test_embeddings.py

@@ -13,10 +13,11 @@
 import sys
 
 import lance
-import lancedb
 import numpy as np
 import pyarrow as pa
 import pytest
+
+import lancedb
 from lancedb.conftest import MockTextEmbeddingFunction
 from lancedb.embeddings import (
     EmbeddingFunctionConfig,

python/python/tests/test_embeddings_slow.py

@@ -14,11 +14,12 @@ import importlib
 import io
 import os
 
-import lancedb
 import numpy as np
 import pandas as pd
 import pytest
 import requests
+
+import lancedb
 from lancedb.embeddings import get_registry
 from lancedb.pydantic import LanceModel, Vector
 
@@ -184,9 +185,10 @@ def test_imagebind(tmp_path):
     import shutil
     import tempfile
 
-    import lancedb.embeddings.imagebind
     import pandas as pd
     import requests
+
+    import lancedb.embeddings.imagebind
     from lancedb.embeddings import get_registry
     from lancedb.pydantic import LanceModel, Vector
 

python/python/tests/test_fts.py

@@ -14,11 +14,12 @@ import os
 import random
 from unittest import mock
 
-import lancedb as ldb
 import numpy as np
 import pandas as pd
 import pytest
 
+import lancedb as ldb
+
 pytest.importorskip("lancedb.fts")
 tantivy = pytest.importorskip("tantivy")
 

python/python/tests/test_io.py

@@ -13,9 +13,10 @@
 
 import os
 
-import lancedb
 import pytest
 
+import lancedb
+
 # You need to setup AWS credentials an a base path to run this test. Example
 #    AWS_PROFILE=default TEST_S3_BASE_URL=s3://my_bucket/dataset pytest tests/test_io.py
 

python/python/tests/test_pydantic.py

@@ -14,18 +14,17 @@
 
 import io
 import json
-import sys
 import os
+import sys
 from datetime import date, datetime
+from pathlib import Path
 from typing import List, Optional, Tuple
 
-from pathlib import Path
 import pyarrow as pa
 import pydantic
 import pytest
-import pytz
-from pydantic import Field
 from lance.arrow import EncodedImageType
+from pydantic import Field
 
 from lancedb.pydantic import (
     PYDANTIC_VERSION,

python/python/tests/test_query.py

@@ -18,6 +18,7 @@ import numpy as np
 import pandas.testing as tm
 import pyarrow as pa
 import pytest
+
 from lancedb.db import LanceDBConnection
 from lancedb.pydantic import LanceModel, Vector
 from lancedb.query import LanceVectorQueryBuilder, Query

python/python/tests/test_remote_client.py

@@ -17,6 +17,7 @@ import pandas as pd
 import pyarrow as pa
 import pytest
 from aiohttp import web
+
 from lancedb.remote.client import RestfulLanceDBClient, VectorQuery
 
 

python/python/tests/test_remote_db.py

@@ -11,8 +11,9 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import lancedb
 import pyarrow as pa
+
+import lancedb
 from lancedb.remote.client import VectorQuery, VectorQueryResult
 
 

python/python/tests/test_rerankers.py

@@ -1,8 +1,9 @@
 import os
 
-import lancedb
 import numpy as np
 import pytest
+
+import lancedb
 from lancedb.conftest import MockTextEmbeddingFunction  # noqa
 from lancedb.embeddings import EmbeddingFunctionRegistry
 from lancedb.pydantic import LanceModel, Vector

python/python/tests/test_table.py

@@ -12,30 +12,31 @@
 #  limitations under the License.
 
 import functools
-from copy import copy
-from datetime import date, datetime, timedelta
 import io
 import os
+from copy import copy
+from datetime import date, datetime, timedelta
 from pathlib import Path
 from time import sleep
 from typing import List
 from unittest.mock import PropertyMock, patch
 
 import lance
-import lancedb
 import numpy as np
 import pandas as pd
 import polars as pl
 import pyarrow as pa
 import pytest
 import pytest_asyncio
+from lance.arrow import EncodedImageType
+from pydantic import BaseModel
+
+import lancedb
 from lancedb.conftest import MockTextEmbeddingFunction
 from lancedb.db import AsyncConnection, LanceDBConnection
 from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry
 from lancedb.pydantic import EncodedImage, LanceModel, Vector
 from lancedb.table import LanceTable
-from pydantic import BaseModel
-from lance.arrow import EncodedImageArray, EncodedImageType, ImageURIType
 
 
 class MockDB:

python/python/tests/test_telemetry.py

@@ -1,7 +1,8 @@
 import json
 
-import lancedb
 import pytest
+
+import lancedb
 from lancedb.utils.events import _Events
 
 

python/python/tests/test_util.py

@@ -15,6 +15,7 @@ import os
 import pathlib
 
 import pytest
+
 from lancedb.util import get_uri_scheme, join_uri