From 832aecaddce80e312095ecdb572d7756eb9df5e9 Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Wed, 9 Sep 2026 20:32:17 -0700 Subject: Ruff auto fix, and requirements bump --- jb/managers/__init__.py | 2 +- jb/managers/amt.py | 28 +++++------ jb/managers/assignment.py | 85 ++++++++++++-------------------- jb/managers/bonus.py | 22 ++++----- jb/managers/email_manager.py | 5 +- jb/managers/gr_api.py | 6 +-- jb/managers/hit.py | 115 +++++++++++++++++-------------------------- jb/managers/thl.py | 19 +++---- jb/managers/worker.py | 6 +-- 9 files changed, 110 insertions(+), 178 deletions(-) (limited to 'jb/managers') diff --git a/jb/managers/__init__.py b/jb/managers/__init__.py index ec64f9f..92ba8bd 100644 --- a/jb/managers/__init__.py +++ b/jb/managers/__init__.py @@ -1,5 +1,5 @@ +from collections.abc import Collection from enum import IntEnum -from typing import Collection from generalresearch.pg_helper import PostgresConfig diff --git a/jb/managers/amt.py b/jb/managers/amt.py index 2cb0cbd..e2c7e90 100644 --- a/jb/managers/amt.py +++ b/jb/managers/amt.py @@ -1,6 +1,6 @@ import logging -from datetime import timezone, datetime -from typing import Tuple, Optional, List, Dict, Any +from datetime import datetime, timezone +from typing import Any import botocore.exceptions from generalresearch.currency import USDCent @@ -9,8 +9,8 @@ from mypy_boto3_mturk.type_defs import ( AssignmentTypeDef, BonusPaymentTypeDef, CreateHITTypeResponseTypeDef, - GetHITResponseTypeDef, CreateHITWithHITTypeResponseTypeDef, + GetHITResponseTypeDef, ) from pydantic import ValidationError @@ -19,7 +19,7 @@ from jb.models import AMTAccount from jb.models.assignment import Assignment from jb.models.bonus import Bonus from jb.models.definitions import HitStatus -from jb.models.hit import HitType, HitQuestion, Hit +from jb.models.hit import Hit, HitQuestion, HitType REJECT_MESSAGE_UNKNOWN_ASSIGNMENT = "Unknown assignment" REJECT_MESSAGE_NO_WORK = "Assignment was submitted with no attempted work." @@ -56,7 +56,7 @@ class AMTManager: } ) - def get_hit_if_exists(self, amt_hit_id: str) -> Tuple[Optional[Hit], Optional[str]]: + def get_hit_if_exists(self, amt_hit_id: str) -> tuple[Hit | None, str | None]: try: res: GetHITResponseTypeDef = self.amt_client.get_hit(HITId=amt_hit_id) @@ -146,7 +146,7 @@ class AMTManager: assert assignment.id is None return assignment - def get_assignment_if_exists(self, amt_assignment_id: str) -> Optional[Assignment]: + def get_assignment_if_exists(self, amt_assignment_id: str) -> Assignment | None: expected_err_msg = f"Assignment {amt_assignment_id} does not exist" try: @@ -161,7 +161,7 @@ class AMTManager: def reject_assignment_if_possible( self, amt_assignment_id: str, msg: str = REJECT_MESSAGE_UNKNOWN_ASSIGNMENT - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: # Unclear to me when this would fail try: @@ -178,7 +178,7 @@ class AMTManager: amt_assignment_id: str, msg: str = APPROVAL_MESSAGE, override_rejection: bool = False, - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: # Unclear to me when this would fail try: @@ -207,8 +207,6 @@ class AMTManager: # elif "This HIT is currently in the state 'Reviewing'" in error_msg: # logging.warning(error_msg) - return None - def send_bonus( self, amt_worker_id: str, @@ -216,7 +214,7 @@ class AMTManager: amt_assignment_id: str, reason: str, unique_request_token: str, - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: try: return self.amt_client.send_bonus( WorkerId=amt_worker_id, @@ -230,11 +228,9 @@ class AMTManager: logging.warning(f"{amt_worker_id=} {amt_assignment_id=}, {e}") return None - def get_bonus( - self, amt_assignment_id: str, payout_event_id: str - ) -> Optional[Bonus]: + def get_bonus(self, amt_assignment_id: str, payout_event_id: str) -> Bonus | None: - res: List[BonusPaymentTypeDef] = self.amt_client.list_bonus_payments( + res: list[BonusPaymentTypeDef] = self.amt_client.list_bonus_payments( AssignmentId=amt_assignment_id )["BonusPayments"] @@ -268,5 +264,3 @@ class AMTManager: self.amt_client.update_expiration_for_hit( HITId=hit["HITId"], ExpireAt=now ) - - return None diff --git a/jb/managers/assignment.py b/jb/managers/assignment.py index f6aa2ce..089adb1 100644 --- a/jb/managers/assignment.py +++ b/jb/managers/assignment.py @@ -1,11 +1,10 @@ from datetime import datetime, timezone -from typing import Optional from psycopg import sql from pydantic import NonNegativeInt, PositiveInt from jb.managers import PostgresManager -from jb.models.assignment import AssignmentStub, Assignment +from jb.models.assignment import Assignment, AssignmentStub from jb.models.definitions import AssignmentStatus @@ -14,8 +13,7 @@ class AssignmentManager(PostgresManager): def create_stub(self, stub: AssignmentStub) -> None: assert stub.id is None data = stub.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_assignment (amt_assignment_id, amt_worker_id, status, created_at, modified_at, hit_id) @@ -23,16 +21,13 @@ class AssignmentManager(PostgresManager): (%(amt_assignment_id)s, %(amt_worker_id)s, %(status)s, %(created_at)s, %(modified_at)s, %(hit_id)s) RETURNING id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() stub.id = pk - return None def create(self, assignment: Assignment) -> None: # Typically this is NOT used (we'd create the stub when HIT is @@ -42,8 +37,7 @@ class AssignmentManager(PostgresManager): assert assignment.id is None data = assignment.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_assignment (amt_assignment_id, amt_worker_id, status, created_at, modified_at, hit_id, @@ -57,16 +51,13 @@ class AssignmentManager(PostgresManager): %(approval_time)s, %(rejection_time)s, %(requester_feedback)s, %(tsid)s) RETURNING id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() assignment.id = pk - return None def get_stub(self, amt_assignment_id: str) -> AssignmentStub: res = self.pg_config.execute_sql_query( @@ -85,7 +76,7 @@ class AssignmentManager(PostgresManager): assert len(res) == 1 return AssignmentStub.model_validate(res[0]) - def get_stub_if_exists(self, amt_assignment_id: str) -> Optional[AssignmentStub]: + def get_stub_if_exists(self, amt_assignment_id: str) -> AssignmentStub | None: try: return self.get_stub(amt_assignment_id=amt_assignment_id) except AssertionError: @@ -121,8 +112,7 @@ class AssignmentManager(PostgresManager): "amt_assignment_id": assignment.amt_assignment_id, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, auto_approval_time = %(auto_approval_time)s, @@ -130,15 +120,12 @@ class AssignmentManager(PostgresManager): tsid = %(tsid)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def reject(self, assignment: Assignment) -> None: assert assignment.status == AssignmentStatus.Rejected @@ -156,8 +143,7 @@ class AssignmentManager(PostgresManager): "accept_time": assignment.accept_time, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, rejection_time = %(rejection_time)s, @@ -167,15 +153,12 @@ class AssignmentManager(PostgresManager): accept_time = %(accept_time)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def approve(self, assignment: Assignment) -> None: assert assignment.status == AssignmentStatus.Approved @@ -194,8 +177,7 @@ class AssignmentManager(PostgresManager): "accept_time": assignment.accept_time, "modified_at": now, } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_assignment SET submit_time = %(submit_time)s, approval_time = %(approval_time)s, @@ -205,15 +187,12 @@ class AssignmentManager(PostgresManager): accept_time = %(accept_time)s, modified_at = %(modified_at)s WHERE amt_assignment_id = %(amt_assignment_id)s - """ - ) + """) # We force this to fail if the assignment doesn't already exist in the db - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" - conn.commit() - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + assert c.rowcount == 1, f"Expected 1 row, got {c.rowcount}" + conn.commit() def missing_tsid_count( self, amt_worker_id: str, lookback_hrs: PositiveInt = 24 diff --git a/jb/managers/bonus.py b/jb/managers/bonus.py index 89b81f0..15d0e5b 100644 --- a/jb/managers/bonus.py +++ b/jb/managers/bonus.py @@ -1,4 +1,4 @@ -from typing import List, Any +from typing import Any from psycopg import sql @@ -11,8 +11,7 @@ class BonusManager(PostgresManager): def create(self, bonus: Bonus) -> None: assert bonus.id is None data = bonus.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_bonus (payout_event_id, amt_worker_id, amount, grant_time, assignment_id, reason) VALUES ( @@ -29,20 +28,17 @@ class BonusManager(PostgresManager): %(reason)s ) RETURNING id, assignment_id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - res = c.fetchone() - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + res = c.fetchone() + conn.commit() bonus.id = res["id"] # type: ignore bonus.assignment_id = res["assignment_id"] # type: ignore - return None - def filter(self, amt_assignment_id: str) -> List[Bonus]: - res: List[Any] = self.pg_config.execute_sql_query( + def filter(self, amt_assignment_id: str) -> list[Bonus]: + res: list[Any] = self.pg_config.execute_sql_query( """ SELECT mb.*, ma.amt_assignment_id FROM mtwerk_bonus mb diff --git a/jb/managers/email_manager.py b/jb/managers/email_manager.py index dcbe167..e740e20 100644 --- a/jb/managers/email_manager.py +++ b/jb/managers/email_manager.py @@ -21,7 +21,7 @@ def get_or_create_contact(email: str, amt_worker_id: str | None = None): return contact_id -def send_login_email_from_url(mautic_url, magic_link) -> None: +def send_login_email_from_url(mautic_url: str, magic_link: str) -> None: email_tokens = { "magic_link": magic_link, } @@ -48,7 +48,8 @@ def send_login_email(email: str, magic_token: str): def send_amt_link_email(email: str, magic_token: str): - # don't actually associate the email with the worker ID until they click the link + # Don't actually associate the email with the worker ID + # until they click the link contact_id = get_or_create_contact(email=email) mautic_url = ( f"{MAUTIC_BASE_URL}/api/emails/{EMAIL_TEMPLATE_ID}/contact/{contact_id}/send" diff --git a/jb/managers/gr_api.py b/jb/managers/gr_api.py index 20dc051..494ceea 100644 --- a/jb/managers/gr_api.py +++ b/jb/managers/gr_api.py @@ -66,7 +66,7 @@ class GRApiManager: "product_user_id": res["product_user_id"], "email": res["metadata"].get("email_address"), "display_name": res["metadata"].get("display_name"), - 'blocked': res['blocked'], + "blocked": res["blocked"], } ) @@ -122,9 +122,7 @@ class GRApiManager: AMT account to a General Research account.""" url = f"{self.base_url}/{self.product_id}/user/{amt_worker_id}/" try: - self._request( - "PATCH", url, json={"product_user_id": user.product_user_id} - ) + self._request("PATCH", url, json={"product_user_id": user.product_user_id}) except GRApiNotFoundError as exc: raise ValueError(f"User {amt_worker_id} does not exist") from exc except GRApiError as exc: diff --git a/jb/managers/hit.py b/jb/managers/hit.py index 63af2d4..a178d4d 100644 --- a/jb/managers/hit.py +++ b/jb/managers/hit.py @@ -1,11 +1,10 @@ from datetime import datetime, timezone -from typing import Optional, List from psycopg import sql from jb.managers import PostgresManager from jb.models.definitions import HitStatus -from jb.models.hit import HitQuestion, HitType, Hit +from jb.models.hit import Hit, HitQuestion, HitType class HitQuestionManager(PostgresManager): @@ -13,21 +12,17 @@ class HitQuestionManager(PostgresManager): def create(self, question: HitQuestion) -> None: assert question.id is None data = question.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_question (url, height) VALUES (%(url)s, %(height)s) RETURNING id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() question.id = pk - return None def get_by_id(self, question_id: int) -> HitQuestion: res = self.pg_config.execute_sql_query( @@ -42,7 +37,7 @@ class HitQuestionManager(PostgresManager): assert len(res) == 1 return HitQuestion.model_validate(res[0]) - def get_by_values_if_exists(self, url: str, height: int) -> Optional[HitQuestion]: + def get_by_values_if_exists(self, url: str, height: int) -> HitQuestion | None: res = self.pg_config.execute_sql_query( """ SELECT * @@ -73,8 +68,7 @@ class HitTypeManager(PostgresManager): assert hit_type.amt_hit_type_id is not None data = hit_type.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_hittype ( amt_hit_type_id, title, @@ -96,27 +90,21 @@ class HitTypeManager(PostgresManager): %(min_active)s ) RETURNING id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() hit_type.id = pk - return None - - def filter_active(self) -> List[HitType]: - res = self.pg_config.execute_sql_query( - """ + def filter_active(self) -> list[HitType]: + res = self.pg_config.execute_sql_query(""" SELECT * FROM mtwerk_hittype WHERE min_active > 0 LIMIT 50 - """ - ) + """) if len(res) == 50: raise ValueError("Too many HitTypes!") @@ -135,7 +123,7 @@ class HitTypeManager(PostgresManager): assert len(res) == 1 return HitType.from_postgres(res[0]) - def get_if_exists(self, amt_hit_type_id: str) -> Optional[HitType]: + def get_if_exists(self, amt_hit_type_id: str) -> HitType | None: try: return self.get(amt_hit_type_id=amt_hit_type_id) except AssertionError: @@ -151,22 +139,18 @@ class HitTypeManager(PostgresManager): def set_min_active(self, hit_type: HitType) -> None: assert hit_type.id, "must be in the db first!" - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_hittype SET min_active = %(min_active)s WHERE id = %(id)s - """ - ) + """) data = {"id": hit_type.id, "min_active": hit_type.min_active} - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - conn.commit() - row_cnt = c.rowcount + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + conn.commit() + row_cnt = c.rowcount assert row_cnt == 1, f"Expected 1 row updated, got {row_cnt}" - return None class HitManager(PostgresManager): @@ -175,8 +159,7 @@ class HitManager(PostgresManager): assert hit.amt_hit_id is not None assert hit.id is None data = hit.to_postgres() - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO mtwerk_hit ( amt_hit_id, hit_type_id, @@ -210,38 +193,32 @@ class HitManager(PostgresManager): %(assignment_available_count)s ) RETURNING id; - """ - ) + """) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() hit.id = pk return hit def update_status(self, amt_hit_id: str, hit_status: HitStatus): now = datetime.now(tz=timezone.utc) - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_hit SET status = %(status)s, modified_at = %(modified_at)s WHERE amt_hit_id = %(amt_hit_id)s; - """ - ) + """) data = { "amt_hit_id": amt_hit_id, "status": hit_status.value, "modified_at": now, } - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - conn.commit() - assert c.rowcount == 1, c.rowcount - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + conn.commit() + assert c.rowcount == 1, c.rowcount def update_hit(self, hit: Hit): hit.modified_at = datetime.now(tz=timezone.utc) @@ -255,8 +232,7 @@ class HitManager(PostgresManager): "amt_hit_id", "modified_at", } - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE mtwerk_hit SET status = %(status)s, review_status = %(review_status)s, assignment_pending_count = %(assignment_pending_count)s, @@ -265,17 +241,14 @@ class HitManager(PostgresManager): modified_at = %(modified_at)s WHERE amt_hit_id = %(amt_hit_id)s RETURNING id; - """ - ) + """) data = hit.model_dump(mode="json", include=fields) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - conn.commit() - assert c.rowcount == 1, c.rowcount - hit.id = c.fetchone()["id"] # type: ignore - return None + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + conn.commit() + assert c.rowcount == 1, c.rowcount + hit.id = c.fetchone()["id"] # type: ignore def get_from_amt_id(self, amt_hit_id: str) -> Hit: res = self.pg_config.execute_sql_query( @@ -324,7 +297,7 @@ class HitManager(PostgresManager): return Hit.from_postgres(res) - def get_from_amt_id_if_exists(self, amt_hit_id: str) -> Optional[Hit]: + def get_from_amt_id_if_exists(self, amt_hit_id: str) -> Hit | None: try: return self.get_from_amt_id(amt_hit_id=amt_hit_id) diff --git a/jb/managers/thl.py b/jb/managers/thl.py index e50fe76..85e0697 100644 --- a/jb/managers/thl.py +++ b/jb/managers/thl.py @@ -1,25 +1,18 @@ +import requests +from generalresearch.currency import USDCent +from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.payout import UserPayoutEvent from generalresearch.models.thl.task_status import TaskStatusResponse from generalresearch.models.thl.wallet.cashout_method import ( - CashoutRequestResponse, CashoutRequestInfo, + CashoutRequestResponse, ) -from generalresearch.models.thl.user_profile import UserProfile -from generalresearch.currency import USDCent - from jb.config import settings - -from generalresearch.models.thl.definitions import PayoutStatus - - -from typing import Optional -import requests - from jb.models.auth import User -def get_task_status(tsid: str) -> Optional[TaskStatusResponse]: +def get_task_status(tsid: str) -> TaskStatusResponse | None: url = f"{settings.fsb_host}{settings.product_id}/status/{tsid}/" d = requests.get(url).json() if d.get("msg") == "invalid tsid": @@ -67,7 +60,7 @@ def get_wallet_balance(amt_worker_id: str) -> USDCent: return USDCent(requests.get(url, params=params).json()["wallet"]["amount"]) -def get_wallet_balance_if_non_negative(amt_worker_id: str) -> Optional[USDCent]: +def get_wallet_balance_if_non_negative(amt_worker_id: str) -> USDCent | None: url = f"{settings.fsb_host}{settings.product_id}/wallet/" params = {"bpuid": amt_worker_id} amt = requests.get(url, params=params).json()["wallet"]["amount"] diff --git a/jb/managers/worker.py b/jb/managers/worker.py index e2d7237..9e7d5e3 100644 --- a/jb/managers/worker.py +++ b/jb/managers/worker.py @@ -1,5 +1,3 @@ -from typing import List - from mypy_boto3_mturk.type_defs import WorkerBlockTypeDef from jb.decorators import AMT_CLIENT @@ -8,9 +6,9 @@ from jb.decorators import AMT_CLIENT class WorkerManager: @staticmethod - def fetch_worker_blocks() -> List[WorkerBlockTypeDef]: + def fetch_worker_blocks() -> list[WorkerBlockTypeDef]: p = AMT_CLIENT.get_paginator("list_worker_blocks") - res: List[WorkerBlockTypeDef] = [] + res: list[WorkerBlockTypeDef] = [] for item in p.paginate(): res.extend(item["WorkerBlocks"]) return res -- cgit v1.2.3