fix(documentai-toolbox): resolve mypy errors

Author: chalmerlowe

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/__init__.py

@@ -25,7 +25,7 @@
 from .utilities import docai_utilities, gcs_utilities
 from .wrappers import document, entity, page
 
-__all__ = (document, page, entity, converter, docai_utilities, gcs_utilities)
+__all__ = ("document", "page", "entity", "converter", "docai_utilities", "gcs_utilities")
 
 
 class Python37DeprecationWarning(DeprecationWarning):  # pragma: NO COVER

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/converters/config/bbox_conversion.py

@@ -190,7 +190,7 @@ def convert_bbox_to_docproto_bbox(block: Block) -> documentai.BoundingPoly:
     y_multiplier = 1.0
     normalized_vertices: List[documentai.NormalizedVertex] = []
 
-    if block.page_width and block.page_height:
+    if block.page_width and block.page_height and block.docproto_width is not None and block.docproto_height is not None:
         x_multiplier = _get_multiplier(
             docproto_coordinate=block.docproto_width,
             external_coordinate=block.page_width,

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/converters/config/block.py

@@ -17,12 +17,12 @@
 import dataclasses
 import json
 from types import SimpleNamespace
-from typing import List, Optional, Type
+from typing import Any, List, Optional, Type, cast
 
 from google.cloud import documentai
 
 
-def _get_target_object(json_data: any, target_object: str) -> Optional[SimpleNamespace]:
+def _get_target_object(json_data: Any, target_object: str) -> Any:
     r"""Returns SimpleNamespace of target_object.
 
     Args:
@@ -72,45 +72,45 @@ class Block:
         page_number:
             Optional.
     """
-    type_: SimpleNamespace = dataclasses.field(init=True, repr=False)
-    text: SimpleNamespace = dataclasses.field(init=True, repr=False)
-    bounding_box: Optional[SimpleNamespace] = dataclasses.field(
+    type_: Any = dataclasses.field(init=True, repr=False)
+    text: Any = dataclasses.field(init=True, repr=False)
+    bounding_box: Any = dataclasses.field(
         init=True, repr=False, default=None
     )
-    block_references: Optional[SimpleNamespace] = dataclasses.field(
+    block_references: Any = dataclasses.field(
         init=True, repr=False, default=None
     )
-    block_id: Optional[SimpleNamespace] = dataclasses.field(
+    block_id: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    confidence: Optional[SimpleNamespace] = dataclasses.field(
+    confidence: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    page_number: Optional[SimpleNamespace] = dataclasses.field(
+    page_number: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    page_width: Optional[SimpleNamespace] = dataclasses.field(
+    page_width: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    page_height: Optional[SimpleNamespace] = dataclasses.field(
+    page_height: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_width: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_width: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_height: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_height: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_type: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_type: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_unit: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_unit: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_x: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_x: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
-    bounding_y: Optional[SimpleNamespace] = dataclasses.field(
+    bounding_y: Any = dataclasses.field(
         init=False, repr=False, default=None
     )
     docproto_width: Optional[float] = dataclasses.field(
@@ -180,7 +180,7 @@ def load_blocks_from_schema(
 
         blocks: List[Block] = []
         ens = _get_target_object(objects, entities)
-        for i in ens:
+        for i in cast(Any, ens):
             entity = i
 
             block_text = ""
@@ -203,11 +203,11 @@ def load_blocks_from_schema(
             b = Block(
                 type_=block_type,
                 text=block_text,
-                bounding_box=_get_target_object(entity, normalized_vertices),
+                bounding_box=_get_target_object(entity, normalized_vertices) if normalized_vertices is not None else None,
             )
 
             if id_:
-                b.id_ = _get_target_object(entity, id_)
+                b.block_id = _get_target_object(entity, id_)
             if confidence:
                 b.confidence = _get_target_object(entity, confidence)
             if page_number and page_number in entity:

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/converters/converter.py

@@ -424,7 +424,7 @@ def convert_from_config(
 
     print("-------- Converting Started --------")
     files, labels, did_not_convert = _get_docproto_files(
-        futures_list, project_id, location, processor_id
+        list(futures_list), project_id, location, processor_id
     )
 
     print("-------- Finished Converting --------")

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/converters/vision_helpers.py

@@ -243,17 +243,17 @@ def _generate_entity_annotations(
     """
     entity_annotations: List[EntityAnnotation] = []
     for token in page_info.page.tokens:
-        v: vision.Vertex = []
+        v: list[vision.Vertex] = []
         if token.layout.bounding_poly.vertices:
             for vertex in token.layout.bounding_poly.vertices:
-                v.append({"x": int(vertex.x), "y": int(vertex.y)})
+                v.append(vision.Vertex(x=int(vertex.x), y=int(vertex.y)))
         else:
             for normalized_vertex in token.layout.bounding_poly.normalized_vertices:
                 v.append(
-                    {
-                        "x": int(normalized_vertex.x * page_info.page.dimension.width),
-                        "y": int(normalized_vertex.y * page_info.page.dimension.height),
-                    }
+                    vision.Vertex(
+                        x=int(normalized_vertex.x * page_info.page.dimension.width),
+                        y=int(normalized_vertex.y * page_info.page.dimension.height),
+                    )
                 )
 
         text_start_index = token.layout.text_anchor.text_segments[0].start_index

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/utilities/gcs_utilities.py

@@ -20,7 +20,7 @@
 
 from google.api_core.gapic_v1 import client_info
 
-from google.cloud import documentai, documentai_toolbox, storage
+from google.cloud import documentai, documentai_toolbox, storage  # type: ignore[attr-defined]
 from google.cloud.documentai_toolbox import constants
 
 
@@ -91,6 +91,7 @@ def get_blobs(
     if gcs_uri:
         gcs_bucket_name, gcs_prefix = split_gcs_uri(gcs_uri)
 
+    assert gcs_prefix is not None
     if re.match(constants.FILE_CHECK_REGEX, gcs_prefix):
         raise ValueError("gcs_prefix cannot contain file types")
 

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/wrappers/document.py

@@ -22,7 +22,7 @@
 import glob
 import os
 import re
-from typing import Dict, Iterator, List, Optional, Type, Union
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Type, Union
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.operation import from_gapic as operation_from_gapic
@@ -51,7 +51,7 @@ def _document_layout_blocks_from_shards(
     shards: List[documentai.Document],
 ) -> Iterator[documentai.Document.DocumentLayout.DocumentLayoutBlock]:
     def extract_blocks(
-        blocks: List[documentai.Document.DocumentLayout.DocumentLayoutBlock],
+        blocks: Iterable[documentai.Document.DocumentLayout.DocumentLayoutBlock],
     ) -> Iterator[documentai.Document.DocumentLayout.DocumentLayoutBlock]:
         queue = collections.deque(blocks)
 
@@ -325,8 +325,9 @@ def _dict_to_bigquery(
     bq_client = bigquery.Client(
         project=project_id, client_info=gcs_utilities._get_client_info()
     )
+    resolved_project_id = project_id or bq_client.project
     table_ref = bigquery.DatasetReference(
-        project=project_id, dataset_id=dataset_name
+        project=resolved_project_id, dataset_id=dataset_name
     ).table(table_name)
 
     job_config = bigquery.LoadJobConfig(
@@ -345,7 +346,7 @@ def _dict_to_bigquery(
 
 
 def _apply_text_offset(
-    documentai_object: Union[Dict[str, Dict], List], text_offset: int
+    documentai_object: Union[Dict[str, Any], List[Any]], text_offset: int
 ) -> None:
     r"""Applies a text offset to all text_segments in `documentai_object`.
 

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/wrappers/entity.py

@@ -63,7 +63,10 @@ class Entity:
 
     _image: Optional[Image.Image] = dataclasses.field(init=False, default=None)
 
-    def __post_init__(self, page_offset: int) -> None:
+    def __post_init__(self, page_offset: Optional[int]) -> None:
+        if page_offset is None:
+            page_offset = 0
+            
         self.type_ = self.documentai_object.type_
 
         if self.documentai_object.mention_text:

packages/google-cloud-documentai-toolbox/google/cloud/documentai_toolbox/wrappers/page.py

@@ -18,7 +18,9 @@
 from abc import ABC
 import dataclasses
 from functools import cached_property
-from typing import Iterable, List, Optional, Type
+from typing import Iterable, List, Optional, Type, TypeVar
+
+T = TypeVar("T", bound="_BasePageElement")
 
 import pandas as pd
 
@@ -181,8 +183,8 @@ def _text_segment(self) -> documentai.Document.TextAnchor.TextSegment:
         return self.documentai_object.layout.text_anchor.text_segments[0]
 
     def _get_children_of_element(
-        self, potential_children: List["_BasePageElement"]
-    ) -> List["_BasePageElement"]:
+        self, potential_children: List[T]
+    ) -> List[T]:
         """
         Filters potential child elements to identify only those fully contained within this element.