chore: typechecking

Author: smotornyuk

ckanext/__init__.py

@@ -2,7 +2,7 @@
 try:
     import pkg_resources
 
-    pkg_resources.declare_namespace(__name__)
+    pkg_resources.declare_namespace(__name__)  # pyright: ignore[reportAttributeAccessIssue]
 except ImportError:
     import pkgutil
 

ckanext/relationship/helpers.py

@@ -23,10 +23,10 @@ def relationship_get_entity_list(
     entity: str, entity_type: str, include_private: bool = True
 ) -> list[dict[str, str]]:
     """Return ids list of specified entity (entity, entity_type)."""
-    context = {}
+
     if entity == "package":
         entity_list = tk.get_action("package_search")(
-            context,
+            {},
             {
                 "fq": f"type:{entity_type}",
                 "rows": 1000,
@@ -36,7 +36,7 @@ def relationship_get_entity_list(
         entity_list = entity_list["results"]
     else:
         entity_list = tk.get_action("relationship_get_entity_list")(
-            context,
+            {},
             {"entity": entity, "entity_type": entity_type},
         )
         entity_list = [
@@ -57,8 +57,8 @@ def relationship_get_current_relations_list(
     related_entity_type = data["related_entity_type"]
     relation_type = data["relation_type"]
 
-    current_relation_by_id = []
-    current_relation_by_name = []
+    current_relation_by_id: list[str] = []
+    current_relation_by_name: list[str] = []
 
     if subject_id:
         current_relation_by_id = tk.get_action("relationship_relations_ids_list")(
@@ -120,13 +120,13 @@ def relationship_get_selected_json(selected_ids: list[str] | None = None) -> str
 def relationship_get_choices_for_related_entity_field(
     field: dict[str, Any],
     current_entity_id: str | None,
-) -> list[str | None]:
+) -> list[tuple[str, str | None]]:
     entities = relationship_get_entity_list(
         field["related_entity"],
         field["related_entity_type"],
     )
 
-    choices = []
+    choices: list[tuple[str, str | None]] = []
 
     for entity in entities:
         if entity["id"] == current_entity_id:
@@ -148,11 +148,11 @@ def relationship_get_choices_for_related_entity_field(
 
         choices.append((entity["id"], entity.get("title") or entity.get("name")))
 
-    choices.sort(key=lambda x: x[1])
+    choices.sort(key=lambda x: x[1])  # pyright: ignore[reportArgumentType, reportCallIssue]
     return choices
 
 
-def relationship_format_autocomplete(packages: dict[str, Any]) -> dict[str, Any]:
+def relationship_format_autocomplete(packages: list[dict[str, Any]]) -> dict[str, Any]:
     return {
         "ResultSet": {
             "Result": [

ckanext/relationship/logic/action.py

@@ -4,7 +4,7 @@
 
 from flask import jsonify
 from flask.wrappers import Response
-from sqlalchemy import or_
+import sqlalchemy as sa
 
 import ckan.plugins.toolkit as tk
 from ckan import authz, logic
@@ -35,15 +35,15 @@ def get_actions():
 @validate(schema.relation_create)
 def relationship_relation_create(
     context: Context, data_dict: dict[str, Any]
-) -> list[dict[str, str]]:
+) -> list[dict[str, Any]]:
     """Create relation with specified type (relation_type) between two entities
     specified by ids (subject_id, object_id). Also create reverse relation.
     """
     tk.check_access("relationship_relation_create", context, data_dict)
 
     subject_id = data_dict["subject_id"]
     object_id = data_dict["object_id"]
-    relation_type = data_dict.get("relation_type")
+    relation_type = data_dict["relation_type"]
     extras = data_dict.get("extras", {})
 
     if Relationship.by_object_id(subject_id, object_id, relation_type):
@@ -89,11 +89,11 @@ def relationship_relation_delete(
         context["session"]
         .query(Relationship)
         .filter(
-            or_(
+            sa.or_(
                 Relationship.subject_id == subject_id,
                 Relationship.subject_id == subject_name,
             ),
-            or_(
+            sa.or_(
                 Relationship.object_id == object_id,
                 Relationship.object_id == object_name,
             ),
@@ -109,11 +109,11 @@ def relationship_relation_delete(
         context["session"]
         .query(Relationship)
         .filter(
-            or_(
+            sa.or_(
                 Relationship.subject_id == object_id,
                 Relationship.subject_id == object_name,
             ),
-            or_(
+            sa.or_(
                 Relationship.object_id == subject_id,
                 Relationship.object_id == subject_name,
             ),
@@ -202,7 +202,7 @@ def relationship_get_entity_list(
 
 @validate(schema.autocomplete)
 def relationship_autocomplete(context: Context, data_dict: dict[str, Any]) -> Response:
-    fq = f'type:{data_dict["entity_type"]} -id:{data_dict["current_entity_id"]}'
+    fq = f"type:{data_dict['entity_type']} -id:{data_dict['current_entity_id']}"
 
     if data_dict.get("owned_only") and not (
         authz.is_sysadmin(tk.current_user.id) and not data_dict.get("check_sysadmin")
@@ -228,9 +228,9 @@ def relationship_autocomplete(context: Context, data_dict: dict[str, Any]) -> Re
             if tk.h.check_access("package_update", {"id": pkg["id"]})
         ]
 
-    format_autocomplete_helper = getattr(
+    format_autocomplete_helper: Any = getattr(
         tk.h,
-        data_dict.get("format_autocomplete_helper"),
+        data_dict.get("format_autocomplete_helper", "relationship_format_autocomplete"),
         tk.h.relationship_format_autocomplete,
     )
 

ckanext/relationship/logic/schema.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
+
 from ckan.logic.schema import validator_args
 from ckan.types import Schema, Validator, ValidatorFactory
 
 
 @validator_args
 def relation_create(
     not_empty: Validator,
-    one_of: Validator,
+    one_of: ValidatorFactory,
     default: ValidatorFactory,
     convert_to_json_if_string: Validator,
     dict_only: Validator,
@@ -26,7 +28,7 @@ def relation_create(
 
 @validator_args
 def relation_delete(
-    not_empty: Validator, one_of: Validator, ignore_missing: Validator
+    not_empty: Validator, one_of: ValidatorFactory, ignore_missing: Validator
 ) -> Schema:
     return {
         "subject_id": [
@@ -44,7 +46,7 @@ def relation_delete(
 
 @validator_args
 def relations_list(
-    not_empty: Validator, one_of: Validator, ignore_missing: Validator
+    not_empty: Validator, one_of: ValidatorFactory, ignore_missing: Validator
 ) -> Schema:
     return {
         "subject_id": [
@@ -66,7 +68,7 @@ def relations_list(
 
 @validator_args
 def relations_ids_list(
-    not_empty: Validator, one_of: Validator, ignore_missing: Validator
+    not_empty: Validator, one_of: ValidatorFactory, ignore_missing: Validator
 ) -> Schema:
     return {
         "subject_id": [
@@ -87,7 +89,7 @@ def relations_ids_list(
 
 
 @validator_args
-def get_entity_list(not_empty: Validator, one_of: Validator) -> Schema:
+def get_entity_list(not_empty: Validator, one_of: ValidatorFactory) -> Schema:
     return {
         "entity": [
             not_empty,

ckanext/relationship/logic/validators.py

@@ -1,16 +1,14 @@
 from __future__ import annotations
 
 import json
-from typing import Any
-
-from six import string_types
+from typing import Any, cast
 
 import ckan.plugins.toolkit as tk
 from ckan.types import Context, FlattenDataDict, FlattenErrorDict, FlattenKey
 
 from ckanext.scheming.validation import (
-    scheming_multiple_choice_output,
-    scheming_validator,
+    scheming_multiple_choice_output,  # pyright: ignore[reportUnknownVariableType]
+    scheming_validator,  # pyright: ignore[reportUnknownVariableType]
 )
 
 
@@ -61,9 +59,9 @@ def validator(
 
 def get_current_relations(
     entity_id: str | None,
-    related_entity: str,
-    related_entity_type: str,
-    relation_type: str,
+    related_entity: str | None,
+    related_entity_type: str | None,
+    relation_type: str | None,
 ):
     if entity_id:
         current_relations = tk.get_action("relationship_relations_list")(
@@ -81,19 +79,24 @@ def get_current_relations(
     return set(current_relations)
 
 
-def get_selected_relations(selected_relations: list | str | None) -> set[str]:
-    if isinstance(selected_relations, string_types) and "," in selected_relations:
+def get_selected_relations(selected_relations: list[Any] | str | None) -> set[str]:
+    if selected_relations is None:
+        selected_relations = []
+
+    if isinstance(selected_relations, str) and "," in selected_relations:
         selected_relations = selected_relations.split(",")
 
     if (
         len(selected_relations) == 1
-        and isinstance(selected_relations[0], string_types)
+        and isinstance(selected_relations[0], str)
         and "," in selected_relations[0]
     ):
         selected_relations = selected_relations[0].split(",")
 
     if selected_relations is not tk.missing:
-        selected_relations = scheming_multiple_choice_output(selected_relations)
+        selected_relations = cast(
+            "list[str]", scheming_multiple_choice_output(selected_relations)
+        )
         selected_relations = [] if selected_relations == [""] else selected_relations
     else:
         selected_relations = []

ckanext/relationship/model/base.py

@@ -1,3 +1,12 @@
+from __future__ import annotations
+
+from typing import Any
 from sqlalchemy.ext.declarative import declarative_base
+import ckan.plugins.toolkit as tk
+
+Base: Any
 
-Base = declarative_base()
+if hasattr(tk, "BaseModel"):
+    Base = tk.BaseModel
+else:
+    Base = declarative_base()

ckanext/relationship/model/relationship.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
+from typing import Any
+from typing_extensions import override
 from datetime import datetime
 
-from sqlalchemy import Column, DateTime, Text, or_
+import sqlalchemy as sa
 from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import Mapped
 
 from ckan import logic, model
 from ckan.model.types import make_uuid
@@ -12,20 +15,31 @@
 
 
 class Relationship(Base):
-    __tablename__ = "relationship_relationship"
-    id: str = Column(Text, primary_key=True, default=make_uuid)
-    subject_id: str = Column(Text, nullable=False)
-    object_id: str = Column(Text, nullable=False)
-    relation_type: str = Column(Text, nullable=False)
-    created_at: datetime = Column(DateTime, nullable=False, default=datetime.utcnow)
-    extras: dict = Column(JSONB, nullable=False, default=dict)
-
-    reverse_relation_type = {
+    __table__: sa.Table = sa.Table(
+        "relationship_relationship",
+        Base.metadata,
+        sa.Column("id", sa.Text, primary_key=True, default=make_uuid),
+        sa.Column("subject_id", sa.Text, nullable=False),
+        sa.Column("object_id", sa.Text, nullable=False),
+        sa.Column("relation_type", sa.Text, nullable=False),
+        sa.Column("created_at", sa.DateTime, nullable=False, default=datetime.utcnow),  # pyright: ignore[reportDeprecated]
+        sa.Column("extras", JSONB, nullable=False, default=dict),
+    )
+
+    id: Mapped[str]
+    subject_id: Mapped[str]
+    object_id: Mapped[str]
+    relation_type: Mapped[str]
+    created_at: Mapped[datetime]
+    extras: Mapped[dict[str, Any]]
+
+    reverse_relation_type: dict[str, Any] = {
         "related_to": "related_to",
         "child_of": "parent_of",
         "parent_of": "child_of",
     }
 
+    @override
     def __repr__(self):
         return (
             "Relationship("
@@ -90,7 +104,7 @@ def by_subject_id(
             object_class = logic.model_name_to_class(model, object_entity)
             q = q.join(
                 object_class,
-                or_(
+                sa.or_(
                     cls.object_id == object_class.id,
                     cls.object_id == object_class.name,
                 ),

ckanext/relationship/plugin.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import contextlib
-from typing import Any
+from typing import Any, cast
 
 import ckan.plugins.toolkit as tk
 from ckan import plugins as p
@@ -86,7 +86,9 @@ def after_dataset_delete(self, context: Context, pkg_dict: dict[str, Any]):
     def before_dataset_index(self, pkg_dict: dict[str, Any]):
         pkg_id = pkg_dict["id"]
         pkg_type = pkg_dict["type"]
-        schema = sch.scheming_get_schema("dataset", pkg_type)
+        schema = cast(
+            "dict[str, Any] | None", sch.scheming_get_schema("dataset", pkg_type)
+        )
         if not schema:
             return pkg_dict
         relations_info = utils.get_relations_info(pkg_type)
@@ -113,7 +115,7 @@ def before_dataset_index(self, pkg_dict: dict[str, Any]):
                 related_entity_type,
                 relation_type,
             )
-            pkg_dict[f'vocab_{field["field_name"]}'] = relations_ids
+            pkg_dict[f"vocab_{field['field_name']}"] = relations_ids
 
             pkg_dict.pop(field["field_name"], None)