chore: minor refactoring for readability/simplicity (#42)

- Fixed Type mismatch on `_table_wrapper_from_documentai_table`
- Simplified boolean checks
- Removed Extra local variables
- Updated function names/docs for clarity

Co-authored-by: Gal Zahavi <38544478+galz10@users.noreply.github.com>

Author: holtskinner

packages/google-cloud-documentai-toolbox/.gitignore

@@ -29,6 +29,7 @@ pip-log.txt
 .nox
 .cache
 .pytest_cache
+.mypy_cache/
 
 
 # Mac

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

@@ -250,7 +250,7 @@ def from_document_path(
 
         Args:
             document_path (str):
-                Required. The path to the resp.
+                Required. The path to the document.json file.
         Returns:
             Document:
                 A document from local document_path.
@@ -316,22 +316,16 @@ def search_pages(
                 A list of Pages.
 
         """
-        if (target_string is None and pattern is None) or (
-            target_string is not None and pattern is not None
-        ):
+        if (target_string and pattern) or (not target_string and not pattern):
             raise ValueError(
                 "Exactly one of target_string and pattern must be specified."
             )
 
         found_pages = []
         for page in self.pages:
             for paragraph in page.paragraphs:
-                if target_string is not None and target_string in paragraph.text:
-                    found_pages.append(page)
-                    break
-                elif (
-                    pattern is not None
-                    and re.search(pattern, paragraph.text) is not None
+                if (target_string and target_string in paragraph.text) or (
+                    pattern and re.search(pattern, paragraph.text)
                 ):
                     found_pages.append(page)
                     break

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

@@ -54,18 +54,14 @@ def to_dataframe(self) -> pd.DataFrame:
                 The DataFrame of the table.
 
         """
-        dataframe = None
-
         if not self.body_rows:
-            dataframe = pd.DataFrame(columns=self.header_rows)
-        else:
-            if self.header_rows != []:
-                dataframe = pd.DataFrame(self.body_rows)
-                dataframe.columns = self.header_rows
-            else:
+            return pd.DataFrame(columns=self.header_rows)
 
-                dataframe = pd.DataFrame(self.body_rows)
-                dataframe.columns = [None] * len(self.body_rows[0])
+        dataframe = pd.DataFrame(self.body_rows)
+        if self.header_rows:
+            dataframe.columns = self.header_rows
+        else:
+            dataframe.columns = [None] * len(self.body_rows[0])
 
         return dataframe
 
@@ -102,13 +98,13 @@ def sample_table_to_csv():
 
 
 def _table_wrapper_from_documentai_table(
-    documentai_table: List[documentai.Document.Page.Table], text: str
+    documentai_table: documentai.Document.Page.Table, text: str
 ) -> Table:
     r"""Returns a Table.
 
     Args:
-        documentai_tables (List[documentai.Document.Page.Table]):
-            Required. A list of documentai.Document.Page.Table.
+        documentai_table (documentai.Document.Page.Table):
+            Required. A documentai.Document.Page.Table.
         text (str):
             Required. UTF-8 encoded text in reading order
             from the document.
@@ -119,22 +115,17 @@ def _table_wrapper_from_documentai_table(
 
     """
 
-    header_rows = []
-    body_rows = []
-
-    header_rows = _table_row_from_documentai_table_row(
+    header_rows = _table_rows_from_documentai_table_rows(
         table_rows=documentai_table.header_rows, text=text
     )
-    body_rows = _table_row_from_documentai_table_row(
+    body_rows = _table_rows_from_documentai_table_rows(
         table_rows=documentai_table.body_rows, text=text
     )
 
-    result = Table(
+    return Table(
         documentai_table=documentai_table, body_rows=body_rows, header_rows=header_rows
     )
 
-    return result
-
 
 @dataclasses.dataclass
 class Paragraph:
@@ -185,13 +176,11 @@ def _text_from_element_with_layout(
 
     result_text = ""
 
-    if element_with_layout.layout.text_anchor.text_segments == []:
+    if not element_with_layout.layout.text_anchor.text_segments:
         return ""
-    else:
-        for text_segment in element_with_layout.layout.text_anchor.text_segments:
-            result_text += text[
-                int(text_segment.start_index) : int(text_segment.end_index)
-            ]
+
+    for text_segment in element_with_layout.layout.text_anchor.text_segments:
+        result_text += text[int(text_segment.start_index) : int(text_segment.end_index)]
 
     return result_text
 
@@ -254,10 +243,10 @@ def _get_lines(lines: List[documentai.Document.Page.Line], text: str) -> List[Li
     return result
 
 
-def _table_row_from_documentai_table_row(
+def _table_rows_from_documentai_table_rows(
     table_rows: List[documentai.Document.Page.Table.TableRow], text: str
 ) -> List[str]:
-    r"""Returns a list rows from table_rows.
+    r"""Returns a list of rows from table_rows.
 
     Args:
         table_rows (List[documentai.Document.Page.Table.TableRow]):

packages/google-cloud-documentai-toolbox/tests/unit/test_page.py

@@ -22,7 +22,9 @@
 
 @pytest.fixture
 def docproto():
-    with open("tests/unit/resources/0/toolbox_invoice_test-0.json", "r") as f:
+    with open(
+        "tests/unit/resources/0/toolbox_invoice_test-0.json", "r", encoding="utf-8"
+    ) as f:
         return documentai.Document.from_json(f.read())
 
 
@@ -117,22 +119,22 @@ def test_table_wrapper_from_documentai_table(docproto):
     assert len(table.header_rows[0]) == 4
 
 
-def test_header_for_table_row_from_documentai_table_row(docproto):
+def test_header_for_table_rows_from_documentai_table_rows(docproto):
     docproto_page = docproto.pages[0]
 
-    header_row = page._table_row_from_documentai_table_row(
+    header_rows = page._table_rows_from_documentai_table_rows(
         table_rows=docproto_page.tables[0].header_rows, text=docproto.text
     )
-    assert header_row == [["Item Description", "Quantity", "Price", "Amount"]]
+    assert header_rows == [["Item Description", "Quantity", "Price", "Amount"]]
 
 
-def test_body_for_table_row_from_documentai_table_row(docproto):
+def test_body_for_table_rows_from_documentai_table_rows(docproto):
     docproto_page = docproto.pages[0]
 
-    body_row = page._table_row_from_documentai_table_row(
+    body_rows = page._table_rows_from_documentai_table_rows(
         table_rows=docproto_page.tables[0].body_rows, text=docproto.text
     )
-    assert body_row == [
+    assert body_rows == [
         ["Tool A", "500", "$1.00", "$500.00"],
         ["Service B", "1", "$900.00", "$900.00"],
         ["Resource C", "50", "$12.00", "$600.00"],