WIP: Enable mypy for code examples

.pre-commit-config.yaml

@@ -29,3 +29,10 @@ repos:
         entry: uv run mypy
         language: system
         files: ^encord/.*py$
+  - repo: local
+    hooks:
+      - id: mypy-docs
+        name: mypy-docs
+        entry: uv run mypy
+        language: system
+        files: ^tests/docs/.*py$

tests/docs/export_labels_all_attributes.py

@@ -5,9 +5,10 @@
 # Import dependencies
 import json
 from collections.abc import Iterable
+from typing import Any, Optional, TypedDict, Union
 
 from encord import EncordUserClient
-from encord.objects import ObjectInstance
+from encord.objects import ObjectInstance, Option
 from encord.objects.attributes import Attribute, ChecklistAttribute, RadioAttribute, TextAttribute
 
 # User input
@@ -41,9 +42,23 @@
 results = []
 
 
+class AnswerData(TypedDict):
+    title: str
+    hash: str
+    nested_attributes: list["AttrData"]
+
+
+class AttrData(TypedDict):
+    frame: int
+    attribute_name: str
+    attribute_hash: str
+    attribute_type: Optional[str]
+    answers: list[Union[AnswerData, dict[str, Any]]]
+
+
 # Function to collect and print attributes
-def extract_and_print_attributes(attribute: Attribute, object_instance: ObjectInstance, frame_number: int):
-    attr_data = {
+def extract_and_print_attributes(attribute: Attribute, object_instance: ObjectInstance, frame_number: int) -> AttrData:
+    attr_data: AttrData = {
         "frame": frame_number + 1,
         "attribute_name": attribute.title,
         "attribute_hash": attribute.feature_node_hash,
@@ -58,28 +73,36 @@ def extract_and_print_attributes(attribute: Attribute, object_instance: ObjectIn
 
     elif isinstance(attribute, RadioAttribute):
         attr_data["attribute_type"] = "RadioAttribute"
-        answer = object_instance.get_answer(attribute)
-        if answer:
-            answer_data = {"title": answer.title, "hash": answer.feature_node_hash, "nested_attributes": []}
-
-            if hasattr(answer, "attributes") and answer.attributes:
-                for nested_attribute in answer.attributes:
-                    nested_result = extract_and_print_attributes(nested_attribute, object_instance, frame_number)
-                    if nested_result:
-                        answer_data["nested_attributes"].append(nested_result)
+        radio_answer = object_instance.get_answer(attribute)
+        assert isinstance(radio_answer, Option)  # RadioAttribute answer is always an Option object
+        if radio_answer:
+            answer_data: AnswerData = {
+                "title": radio_answer.title,
+                "hash": radio_answer.feature_node_hash,
+                "nested_attributes": [],
+            }
+
+            for nested_attribute in radio_answer.attributes:
+                nested_result = extract_and_print_attributes(nested_attribute, object_instance, frame_number)
+                if nested_result:
+                    answer_data["nested_attributes"].append(nested_result)
 
             attr_data["answers"].append(answer_data)
         else:
             attr_data["answers"].append({"note": "No attribute answer"})
 
     elif isinstance(attribute, ChecklistAttribute):
         attr_data["attribute_type"] = "ChecklistAttribute"
-        answers = object_instance.get_answer(attribute)
-        if answers:
-            if not isinstance(answers, Iterable):
-                answers = [answers]
-            for answer in answers:
-                attr_data["answers"].append({"title": answer.title, "hash": answer.feature_node_hash})
+
+        checklist_answers = object_instance.get_answer(attribute)
+        # Answers of checklist attribute is always a list of Option objects
+        assert isinstance(checklist_answers, Iterable)
+        if checklist_answers:
+            for checklist_answer in checklist_answers:
+                assert isinstance(checklist_answer, Option)  # Enforcing type here
+                attr_data["answers"].append(
+                    {"title": checklist_answer.title, "hash": checklist_answer.feature_node_hash}
+                )
         else:
             attr_data["answers"].append({"note": "No attribute answer"})
 
@@ -103,7 +126,7 @@ def extract_and_print_attributes(attribute: Attribute, object_instance: ObjectIn
 # Process all label rows
 for label_row in label_rows:
     object_instances = label_row.get_object_instances()
-    assert object_instances, f"No object instances found in label row {label_row.uid}"
+    assert object_instances, f"No object instances found in label row {label_row.label_hash}"
 
     for object_instance in object_instances:
         annotations = object_instance.get_annotations()

tests/docs/export_labels_text_attributes.py

@@ -56,7 +56,7 @@ def extract_text_attributes(attribute: Attribute, object_instance: ObjectInstanc
 # Iterate through all object instances and collect text attribute data
 for label_row in label_rows:
     object_instances = label_row.get_object_instances()
-    assert object_instances, f"No object instances found in label row {label_row.uid}"
+    assert object_instances, f"No object instances found in label row {label_row.label_hash}"
 
     for object_instance in object_instances:
         annotations = object_instance.get_annotations()

tests/docs/label_bitmasks_images_videos_example.py

@@ -4,6 +4,7 @@
 
 # Import dependencies
 import os
+from typing import Any
 
 import numpy as np
 
@@ -74,7 +75,7 @@ def assert_checklist_and_options(title, *option_titles):
 )
 assert other_apple_option_text_attribute is not None, "TextAttribute 'Specify apple type' not found"
 
-video_frame_labels = {
+video_frame_labels: dict[str, Any] = {
     "cherries-001.jpg": {
         0: {
             "label_ref": "apple_001",
@@ -243,11 +244,12 @@ def assert_checklist_and_options(title, *option_titles):
 label_rows_to_save = []
 
 for data_unit, frame_coordinates in video_frame_labels.items():
-    label_row = label_row_map.get(data_unit)
-    if not label_row:
+    if data_unit not in label_row_map:
         print(f"⚠️ Skipping: No initialized label row found for {data_unit}")
         continue
 
+    label_row = label_row_map[data_unit]
+
     object_instances_by_label_ref = {}
 
     for frame_number, items in frame_coordinates.items():

tests/docs/label_classifications_example.py

@@ -4,6 +4,7 @@
 
 # Import dependencies
 from pathlib import Path
+from typing import Any
 
 from encord import EncordUserClient
 from encord.objects import Classification, Option
@@ -24,7 +25,7 @@
 project = user_client.get_project(PROJECT_ID)
 
 # Define specific configurations for each data unit
-data_unit_configs = [
+data_unit_configs: Any = [
     {
         "title": "anne-of-green-gables.pdf",
         "classifications": [

tests/docs/label_polygons_pdfs_example.py

@@ -151,10 +151,10 @@
             print(f"[SKIP] No label row found for: {data_unit}")
             continue
 
-        label_row = label_rows[0]
+        label_row_to_initialise = label_rows[0]
         try:
-            label_row.initialise_labels(bundle=bundle)
-            label_row_map[data_unit] = label_row
+            label_row_to_initialise.initialise_labels(bundle=bundle)
+            label_row_map[data_unit] = label_row_to_initialise
             print(f"[INIT] Initialized label row for: {data_unit}")
         except Exception as e:
             raise AssertionError(f"[ASSERT] Failed to initialize label row for {data_unit}: {e}")
@@ -168,14 +168,14 @@
     object_instances_by_label_ref = {}
 
     for frame_number, items in frame_coordinates.items():
-        if not isinstance(items, list):
-            items = [items]
-
-        for item in items:
+        items_list = items if isinstance(items, list) else [items]
+        for item in items_list:
             label_ref = item["label_ref"]
             coord = item["coordinates"]
             correction_type = item["correction_type"]
             checklist_options_str = item.get("checklist_options", "")
+            assert isinstance(checklist_options_str, str), "[ASSERT] checklist_options must be a string"
+
             text_correction = item.get("text_correction", "")
 
             # Basic checks

tests/docs/project-timers-stage.py

@@ -28,9 +28,12 @@
 time_entries = list(project.list_time_spent(start=start_date, end=end_date))
 
 # Filters for total time by stage
-stage_time = defaultdict(int)
+stage_time: dict[str, int] = defaultdict(int)
 for entry in time_entries:
-    stage_time[entry.workflow_stage.title] += entry.time_spent_seconds
+    if entry.workflow_stage is None:
+        stage_time["Time outside of the queue"] += entry.time_spent_seconds
+    else:
+        stage_time[entry.workflow_stage.title] += entry.time_spent_seconds
 
 for stage, seconds in sorted(stage_time.items(), key=lambda x: x[1], reverse=True):
     print(f"Stage: {stage:<25} | {seconds:>5} sec | {seconds / 60:>5.1f} min")

tests/docs/project-timers-task.py

@@ -29,7 +29,7 @@
 
 # Filters for task time
 
-task_time = defaultdict(int)
+task_time: dict[str, int] = defaultdict(int)
 for entry in time_entries:
     task_time[str(entry.workflow_task_uuid)] += entry.time_spent_seconds