diff --git a/prefect_gitlab/credentials.py b/prefect_gitlab/credentials.py index fed1829..107ac5f 100644 --- 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 index 1cb9ded..a2b0f1d 100644 --- 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 index a52370d..7dc369b 100644 --- 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