Conditional imports to support pydantic>2

PrefectHQ · Oct 4, 2023 · 47eb029 · 47eb029
1 parent d5f4064
commit 47eb029
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/prefect_gitlab/credentials.py b/prefect_gitlab/credentials.py
@@ -2,7 +2,12 @@
 
 from gitlab import Gitlab
 from prefect.blocks.core import Block
-from pydantic import Field, HttpUrl, SecretStr
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, HttpUrl, SecretStr
+else:
+    from pydantic import Field, HttpUrl, SecretStr
 
 
 class GitLabCredentials(Block):

diff --git a/prefect_gitlab/repositories.py b/prefect_gitlab/repositories.py
@@ -51,7 +51,12 @@
 from prefect.filesystems import ReadableDeploymentStorage
 from prefect.utilities.asyncutils import sync_compatible
 from prefect.utilities.processutils import run_process
-from pydantic import Field, HttpUrl, validator
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, HttpUrl, validator
+else:
+    from pydantic import Field, HttpUrl, validator
 
 from prefect_gitlab.credentials import GitLabCredentials
 

diff --git a/tests/test_repositories.py b/tests/test_repositories.py
@@ -6,7 +6,12 @@
 import pytest
 from prefect.exceptions import InvalidRepositoryURLError
 from prefect.testing.utilities import AsyncMock
-from pydantic import SecretStr
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import SecretStr
+else:
+    from pydantic import SecretStr
 
 import prefect_gitlab
 from prefect_gitlab.credentials import GitLabCredentials