Skip to content

Commit bb46914

Browse files
authored
Merge pull request #161 from enoch3712/139-change-name-documentloaderdocumentai
name convention changed with warning
2 parents 9ac8dab + 4c41f24 commit bb46914

File tree

3 files changed

+39
-4
lines changed

3 files changed

+39
-4
lines changed

extract_thinker/__init__.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@
2222
from .document_loader.document_loader_doc2txt import DocumentLoaderDoc2txt
2323
from .document_loader.document_loader_aws_textract import DocumentLoaderAWSTextract
2424
from .document_loader.document_loader_llm_image import DocumentLoaderLLMImage
25-
from .document_loader.document_loader_google_document_ai import DocumentLoaderDocumentAI
25+
from .document_loader.document_loader_google_document_ai import (
26+
DocumentLoaderGoogleDocumentAI,
27+
DocumentLoaderDocumentAI,
28+
)
2629

2730
__all__ = [
2831
'Extractor',
@@ -39,6 +42,7 @@
3942
'DocumentLoaderTxt',
4043
'DocumentLoaderDoc2txt',
4144
'DocumentLoaderAWSTextract',
45+
'DocumentLoaderGoogleDocumentAI',
4246
'DocumentLoaderDocumentAI',
4347
'Classification',
4448
'ClassificationResponse',

extract_thinker/document_loader/document_loader_google_document_ai.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from cachetools.keys import hashkey
1111
from cachetools import cachedmethod
1212
from operator import attrgetter
13+
import warnings
1314

1415
class Config:
1516
def __init__(
@@ -21,7 +22,7 @@ def __init__(
2122
self.page_range = page_range
2223

2324

24-
class DocumentLoaderDocumentAI(CachedDocumentLoader):
25+
class DocumentLoaderGoogleDocumentAI(CachedDocumentLoader):
2526
"""Loader for documents using Google Document AI."""
2627

2728
SUPPORTED_FORMATS = [
@@ -179,3 +180,14 @@ def _get_page_key_value_pairs(page: documentai.Document.Page) -> List[Dict[str,
179180
for kv_pair in page.key_value_pairs
180181
]
181182

183+
184+
# Create an alias with deprecation warning
185+
class DocumentLoaderDocumentAI(DocumentLoaderGoogleDocumentAI):
186+
def __init__(self, *args, **kwargs):
187+
warnings.warn(
188+
"DocumentLoaderDocumentAI is deprecated and will be removed in 0.1.0"
189+
"Use DocumentLoaderGoogleDocumentAI instead.",
190+
DeprecationWarning,
191+
stacklevel=2
192+
)
193+
super().__init__(*args, **kwargs)

tests/test_document_loader_google_document_ai.py

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,40 @@
11
import os
22
import pytest
3+
import warnings
34
from dotenv import load_dotenv
4-
from extract_thinker.document_loader.document_loader_google_document_ai import DocumentLoaderDocumentAI
5+
from extract_thinker.document_loader.document_loader_google_document_ai import (
6+
DocumentLoaderGoogleDocumentAI,
7+
DocumentLoaderDocumentAI
8+
)
59
from tests.test_document_loader_base import BaseDocumentLoaderTest
610

711
load_dotenv()
812

913
class TestDocumentLoaderGoogleDocumentAI(BaseDocumentLoaderTest):
1014
@pytest.fixture
1115
def loader(self):
12-
return DocumentLoaderDocumentAI(
16+
return DocumentLoaderGoogleDocumentAI(
1317
project_id=os.getenv("DOCUMENTAI_PROJECT_ID"),
1418
location=os.getenv("DOCUMENTAI_LOCATION"),
1519
processor_id=os.getenv("DOCUMENTAI_PROCESSOR_ID"),
1620
credentials=os.getenv("DOCUMENTAI_GOOGLE_CREDENTIALS")
1721
)
1822

23+
def test_deprecation_warning(self):
24+
"""Test that using old class name raises deprecation warning"""
25+
with pytest.warns(DeprecationWarning) as record:
26+
DocumentLoaderDocumentAI(
27+
project_id=os.getenv("DOCUMENTAI_PROJECT_ID"),
28+
location=os.getenv("DOCUMENTAI_LOCATION"),
29+
processor_id=os.getenv("DOCUMENTAI_PROCESSOR_ID"),
30+
credentials=os.getenv("DOCUMENTAI_GOOGLE_CREDENTIALS")
31+
)
32+
33+
# Verify the warning message
34+
assert len(record) == 1
35+
assert "DocumentLoaderDocumentAI is deprecated" in str(record[0].message)
36+
assert "Use DocumentLoaderGoogleDocumentAI instead" in str(record[0].message)
37+
1938
@pytest.fixture
2039
def test_file_path(self):
2140
current_dir = os.path.dirname(os.path.abspath(__file__))

0 commit comments

Comments
 (0)