From cff96d2c7ad155f3c3b5e51c76f8d46696457a09 Mon Sep 17 00:00:00 2001 From: stuppie Date: Fri, 21 Aug 2026 14:55:42 -0600 Subject: fix import (fastapi.Request -> pytest.FixtureRequest) --- test_utils/models/network/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'test_utils/models/network') diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py index abfbc18..bebc691 100644 --- a/test_utils/models/network/conftest.py +++ b/test_utils/models/network/conftest.py @@ -3,7 +3,7 @@ from datetime import datetime, timedelta, timezone from uuid import uuid4 import pytest -from fastapi import Request +from pytest import FixtureRequest as Request from generalresearch.managers.network.label import IPLabelManager from generalresearch.managers.network.tool_run import ToolRunManager -- cgit v1.2.3 From 6cf7ccbaa8306700e64ada19d6f99807743b2865 Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Fri, 21 Aug 2026 17:19:41 -0700 Subject: Ruff auto updates to 3.14 --- generalresearch/__init__.py | 5 +- generalresearch/config.py | 4 +- generalresearch/grliq/managers/event_plotter.py | 6 +- generalresearch/grliq/managers/forensic_data.py | 3 +- generalresearch/grliq/managers/forensic_events.py | 29 +++---- generalresearch/grliq/managers/forensic_results.py | 23 +++--- generalresearch/grliq/managers/forensic_summary.py | 14 ++-- generalresearch/grliq/models/custom_types.py | 2 +- generalresearch/grliq/models/decider.py | 4 +- generalresearch/grliq/models/events.py | 32 ++++---- generalresearch/grliq/models/forensic_data.py | 6 +- generalresearch/grliq/models/forensic_summary.py | 2 +- generalresearch/grliq/models/useragents.py | 2 +- generalresearch/grliq/utils.py | 4 +- generalresearch/grpc.py | 8 +- generalresearch/incite/base.py | 25 +++--- generalresearch/incite/defaults.py | 60 +++++++------- generalresearch/incite/mergers/__init__.py | 24 +++--- generalresearch/incite/mergers/ym_wall_summary.py | 2 +- generalresearch/incite/schemas/__init__.py | 4 +- .../mergers/foundations/enriched_task_adjust.py | 2 +- generalresearch/incite/schemas/thl_web.py | 10 +-- generalresearch/locales/__init__.py | 4 +- generalresearch/locales/timezone.py | 4 +- generalresearch/managers/cint/profiling.py | 2 +- generalresearch/managers/cint/survey.py | 4 +- generalresearch/managers/criteria.py | 4 +- generalresearch/managers/dynata/profiling.py | 2 +- generalresearch/managers/dynata/survey.py | 6 +- generalresearch/managers/events.py | 6 +- generalresearch/managers/gr/authentication.py | 12 +-- generalresearch/managers/gr/team.py | 4 +- generalresearch/managers/innovate/survey.py | 6 +- generalresearch/managers/leaderboard/manager.py | 6 +- generalresearch/managers/morning/survey.py | 8 +- generalresearch/managers/network/label.py | 6 +- generalresearch/managers/precision/survey.py | 6 +- generalresearch/managers/prodege/survey.py | 8 +- generalresearch/managers/repdata/survey.py | 6 +- generalresearch/managers/sago/survey.py | 8 +- generalresearch/managers/spectrum/survey.py | 8 +- generalresearch/managers/thl/buyer.py | 4 +- generalresearch/managers/thl/cashout_method.py | 4 +- generalresearch/managers/thl/contest_manager.py | 8 +- .../managers/thl/ledger_manager/conditions.py | 7 +- .../managers/thl/ledger_manager/ledger.py | 13 ++-- .../managers/thl/ledger_manager/thl_ledger.py | 17 ++-- generalresearch/managers/thl/product.py | 4 +- generalresearch/managers/thl/profiling/uqa.py | 4 +- generalresearch/managers/thl/profiling/user_upk.py | 4 +- generalresearch/managers/thl/session.py | 20 ++--- generalresearch/managers/thl/survey.py | 4 +- generalresearch/managers/thl/task_adjustment.py | 6 +- generalresearch/managers/thl/user_compensate.py | 4 +- .../thl/user_manager/mysql_user_manager.py | 6 +- generalresearch/managers/thl/userhealth.py | 8 +- generalresearch/managers/thl/wall.py | 12 ++- generalresearch/managers/thl/wallet/__init__.py | 4 +- generalresearch/managers/thl/wallet/tango.py | 4 +- generalresearch/models/admin/__init__.py | 14 ++-- generalresearch/models/admin/request.py | 12 +-- generalresearch/models/cint/__init__.py | 3 +- generalresearch/models/cint/question.py | 9 +-- generalresearch/models/cint/survey.py | 13 ++-- generalresearch/models/cint/task_collection.py | 4 +- generalresearch/models/custom_types.py | 15 ++-- generalresearch/models/dynata/survey.py | 11 ++- generalresearch/models/events.py | 70 +++++++---------- generalresearch/models/gr/authentication.py | 11 ++- generalresearch/models/gr/business.py | 11 ++- generalresearch/models/gr/team.py | 7 +- generalresearch/models/innovate/__init__.py | 2 +- generalresearch/models/innovate/survey.py | 14 ++-- generalresearch/models/legacy/bucket.py | 3 +- generalresearch/models/legacy/questions.py | 3 +- generalresearch/models/lucid/__init__.py | 3 +- generalresearch/models/lucid/question.py | 3 +- generalresearch/models/marketplace/summary.py | 3 +- generalresearch/models/morning/__init__.py | 2 +- generalresearch/models/morning/question.py | 13 ++-- generalresearch/models/morning/survey.py | 84 ++++++++++---------- generalresearch/models/network/mtr/execute.py | 6 +- generalresearch/models/network/nmap/parser.py | 6 +- generalresearch/models/network/nmap/result.py | 4 +- generalresearch/models/network/rdns/execute.py | 6 +- generalresearch/models/precision/__init__.py | 2 +- generalresearch/models/precision/survey.py | 69 ++++++++-------- .../models/precision/task_collection.py | 8 +- generalresearch/models/prodege/__init__.py | 3 +- generalresearch/models/prodege/question.py | 6 +- generalresearch/models/prodege/survey.py | 12 +-- generalresearch/models/prodege/task_collection.py | 4 +- generalresearch/models/repdata/survey.py | 13 ++-- generalresearch/models/sago/__init__.py | 2 +- generalresearch/models/sago/survey.py | 13 ++-- generalresearch/models/spectrum/__init__.py | 2 +- generalresearch/models/spectrum/question.py | 15 ++-- generalresearch/models/spectrum/survey.py | 19 ++--- generalresearch/models/string_utils.py | 2 +- generalresearch/models/thl/category.py | 3 +- generalresearch/models/thl/contest/__init__.py | 6 +- generalresearch/models/thl/contest/contest.py | 19 +++-- .../models/thl/contest/contest_entry.py | 10 +-- generalresearch/models/thl/contest/io.py | 4 +- generalresearch/models/thl/contest/leaderboard.py | 9 +-- generalresearch/models/thl/contest/milestone.py | 2 +- generalresearch/models/thl/contest/raffle.py | 7 +- generalresearch/models/thl/finance.py | 8 +- generalresearch/models/thl/ipinfo.py | 13 ++-- generalresearch/models/thl/leaderboard.py | 8 +- generalresearch/models/thl/ledger.py | 21 ++--- generalresearch/models/thl/ledger_example.py | 10 +-- generalresearch/models/thl/offerwall/__init__.py | 3 +- generalresearch/models/thl/offerwall/base.py | 3 +- generalresearch/models/thl/offerwall/cache.py | 10 +-- generalresearch/models/thl/payout_format.py | 2 +- generalresearch/models/thl/product.py | 8 +- .../models/thl/profiling/marketplace.py | 6 +- .../models/thl/profiling/other_option.py | 2 +- .../models/thl/profiling/upk_question.py | 13 ++-- .../models/thl/profiling/upk_question_answer.py | 9 +-- .../models/thl/profiling/user_question_answer.py | 14 ++-- generalresearch/models/thl/session.py | 27 +++---- generalresearch/models/thl/survey/__init__.py | 2 +- generalresearch/models/thl/survey/buyer.py | 4 +- generalresearch/models/thl/survey/condition.py | 3 +- generalresearch/models/thl/survey/model.py | 17 ++-- generalresearch/models/thl/survey/penalty.py | 9 +-- generalresearch/models/thl/task_adjustment.py | 6 +- generalresearch/models/thl/task_status.py | 3 +- generalresearch/models/thl/user.py | 13 ++-- generalresearch/models/thl/user_iphistory.py | 8 +- generalresearch/models/thl/user_profile.py | 3 +- generalresearch/models/thl/user_quality_event.py | 6 +- generalresearch/models/thl/userhealth.py | 21 +++-- .../models/thl/wallet/cashout_method.py | 6 +- generalresearch/models/thl/wallet/payout.py | 8 +- generalresearch/pg_helper.py | 4 +- generalresearch/sql_helper.py | 4 +- generalresearch/utils/aggregation.py | 2 +- generalresearch/utils/copying_cache.py | 2 +- generalresearch/utils/enum.py | 2 +- generalresearch/wall_status_codes/__init__.py | 6 +- generalresearch/wall_status_codes/cint.py | 6 +- generalresearch/wall_status_codes/dynata.py | 16 ++-- generalresearch/wall_status_codes/fullcircle.py | 14 ++-- generalresearch/wall_status_codes/innovate.py | 12 +-- generalresearch/wall_status_codes/lucid.py | 16 ++-- generalresearch/wall_status_codes/morning.py | 14 ++-- generalresearch/wall_status_codes/pollfish.py | 14 ++-- generalresearch/wall_status_codes/precision.py | 12 +-- generalresearch/wall_status_codes/prodege.py | 10 +-- generalresearch/wall_status_codes/repdata.py | 14 ++-- generalresearch/wall_status_codes/sago.py | 16 ++-- generalresearch/wall_status_codes/spectrum.py | 12 +-- generalresearch/wall_status_codes/wxet.py | 14 ++-- generalresearch/wxet/models/definitions.py | 34 ++++---- generalresearch/wxet/models/finish_type.py | 8 +- test_utils/conftest.py | 91 ++++++++++++++-------- test_utils/grliq/conftest.py | 8 +- test_utils/incite/collections/conftest.py | 3 +- test_utils/incite/conftest.py | 7 +- test_utils/incite/mergers/conftest.py | 2 +- test_utils/managers/conftest.py | 2 +- test_utils/managers/gr/conftest.py | 2 +- test_utils/managers/thl/conftest.py | 2 +- test_utils/managers/upk/conftest.py | 2 +- test_utils/models/conftest.py | 7 +- test_utils/models/contest/conftest.py | 7 +- test_utils/models/gr/conftest.py | 2 +- test_utils/models/ledger/conftest.py | 3 +- test_utils/models/network/conftest.py | 6 +- test_utils/models/thl/conftest.py | 17 ++-- test_utils/spectrum/conftest.py | 16 ++-- tests/grliq/models/test_forensic_data.py | 8 +- .../incite/collections/test_df_collection_base.py | 24 +++--- .../collections/test_df_collection_item_base.py | 20 ++--- .../collections/test_df_collection_item_thl_web.py | 10 +-- .../test_df_collection_thl_marketplaces.py | 6 +- .../collections/test_df_collection_thl_web.py | 2 +- .../mergers/foundations/test_enriched_session.py | 8 +- .../mergers/foundations/test_enriched_wall.py | 16 ++-- .../mergers/foundations/test_user_id_product.py | 16 ++-- tests/incite/mergers/test_merge_collection.py | 8 +- tests/incite/mergers/test_pop_ledger.py | 12 +-- tests/incite/mergers/test_ym_survey_merge.py | 18 ++--- tests/incite/schemas/test_admin_responses.py | 29 ++++--- tests/incite/test_collection_base.py | 22 +++--- tests/incite/test_collection_base_item.py | 6 +- tests/managers/leaderboard.py | 14 ++-- tests/managers/test_events.py | 16 ++-- .../managers/thl/test_contest/test_leaderboard.py | 10 ++- tests/managers/thl/test_contest/test_milestone.py | 16 ++-- tests/managers/thl/test_contest/test_raffle.py | 22 +++--- tests/managers/thl/test_harmonized_uqa.py | 12 +-- tests/managers/thl/test_ledger/test_lm_accounts.py | 65 ++++++++-------- tests/managers/thl/test_ledger/test_lm_tx_locks.py | 22 +++--- .../thl/test_ledger/test_thl_lm_bp_payout.py | 48 ++++++------ tests/managers/thl/test_ledger/test_thl_lm_tx.py | 55 +++++++------ .../test_ledger/test_thl_lm_tx__user_payouts.py | 10 +-- tests/managers/thl/test_ledger/test_user_txs.py | 31 ++++---- tests/managers/thl/test_maxmind.py | 7 +- tests/managers/thl/test_profiling/test_user_upk.py | 4 +- tests/managers/thl/test_survey.py | 12 +-- tests/managers/thl/test_task_adjustment.py | 12 ++- tests/managers/thl/test_task_status.py | 16 ++-- tests/managers/thl/test_user_manager/test_base.py | 4 +- tests/managers/thl/test_user_streak.py | 16 ++-- tests/managers/thl/test_userhealth.py | 8 +- tests/managers/thl/test_wall_manager.py | 12 +-- tests/models/admin/test_report_request.py | 20 ++--- tests/models/custom_types/test_aware_datetime.py | 6 +- tests/models/custom_types/test_dsn.py | 6 +- tests/models/dynata/test_eligbility.py | 6 +- tests/models/gr/test_authentication.py | 6 +- tests/models/gr/test_base.py | 24 +++--- tests/models/gr/test_business.py | 20 ++--- tests/models/morning/test.py | 6 +- tests/models/prodege/test_survey_participation.py | 6 +- tests/models/spectrum/test_question.py | 14 ++-- tests/models/spectrum/test_survey.py | 26 +++---- tests/models/spectrum/test_survey_manager.py | 13 ++-- tests/models/test_finance.py | 10 +-- tests/models/thl/test_adjustments.py | 14 ++-- tests/models/thl/test_contest/test_contest.py | 2 +- .../thl/test_contest/test_leaderboard_contest.py | 6 +- tests/models/thl/test_ledger.py | 4 +- tests/models/thl/test_product.py | 14 ++-- tests/models/thl/test_user.py | 40 +++++----- tests/models/thl/test_user_iphistory.py | 4 +- tests/models/thl/test_wall.py | 38 ++++----- tests/models/thl/test_wall_session.py | 20 ++--- tests/test_postgres.py | 4 +- 233 files changed, 1298 insertions(+), 1379 deletions(-) (limited to 'test_utils/models/network') diff --git a/generalresearch/__init__.py b/generalresearch/__init__.py index 2100d41..604b7e2 100644 --- a/generalresearch/__init__.py +++ b/generalresearch/__init__.py @@ -1,7 +1,8 @@ import threading import time +from collections.abc import Callable from functools import wraps -from typing import Any, Callable, Optional +from typing import Any, Optional from decorator import decorator from wrapt import FunctionWrapper, ObjectProxy @@ -12,7 +13,7 @@ def retry( tries: int = 4, delay: float = 0.5, backoff: int = 2, - logger: Optional[Any] = None, + logger: Any | None = None, ) -> Callable: """ https://www.calazan.com/retry-decorator-for-python-3/ diff --git a/generalresearch/config.py b/generalresearch/config.py index 44f3db7..76e3995 100644 --- a/generalresearch/config.py +++ b/generalresearch/config.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from pathlib import Path from pydantic import DirectoryPath, Field, MariaDBDsn, PostgresDsn, RedisDsn @@ -125,4 +125,4 @@ EXAMPLE_PRODUCT_ID = "1108d053e4fa47c5b0dbdcd03a7981e7" # AMT accounting was changed many times and txs before this date # are either missing AMT bonuses, or not accounting for hit rewards. JAMES_BILLINGS_BPID = "888dbc589987425fa846d6e2a8daed04" -JAMES_BILLINGS_TX_CUTOFF = datetime(2026, 1, 1, tzinfo=timezone.utc) +JAMES_BILLINGS_TX_CUTOFF = datetime(2026, 1, 1, tzinfo=UTC) diff --git a/generalresearch/grliq/managers/event_plotter.py b/generalresearch/grliq/managers/event_plotter.py index 54105ce..ed01d2e 100644 --- a/generalresearch/grliq/managers/event_plotter.py +++ b/generalresearch/grliq/managers/event_plotter.py @@ -11,7 +11,7 @@ from generalresearch.grliq.models.events import KeyboardEvent, MouseEvent def make_events_svg( - mouse_events: List[MouseEvent], keyboard_events: List[KeyboardEvent] + mouse_events: list[MouseEvent], keyboard_events: list[KeyboardEvent] ) -> str: if len(mouse_events) + len(keyboard_events) == 0: return f'\n' + "\n" @@ -119,8 +119,8 @@ def svg_multiline_text( def group_input_events_by_xy( - mouse_events: List[MouseEvent], keyboard_events: List[KeyboardEvent] -) -> List[tuple[tuple[float, float], List[str]]]: + mouse_events: list[MouseEvent], keyboard_events: list[KeyboardEvent] +) -> list[tuple[tuple[float, float], list[str]]]: """ Each keypress is its own event. For plotting, we want to group together all keypresses that were made when the mouse was at the same position, diff --git a/generalresearch/grliq/managers/forensic_data.py b/generalresearch/grliq/managers/forensic_data.py index 739c520..7567552 100644 --- a/generalresearch/grliq/managers/forensic_data.py +++ b/generalresearch/grliq/managers/forensic_data.py @@ -1,7 +1,8 @@ from __future__ import annotations from datetime import datetime -from typing import Any, Collection +from typing import Any +from collections.abc import Collection from psycopg import sql from pydantic import NonNegativeInt, PositiveInt diff --git a/generalresearch/grliq/managers/forensic_events.py b/generalresearch/grliq/managers/forensic_events.py index bbc6b6d..85e9620 100644 --- a/generalresearch/grliq/managers/forensic_events.py +++ b/generalresearch/grliq/managers/forensic_events.py @@ -1,6 +1,7 @@ import json from datetime import datetime -from typing import Any, Collection, Dict, List, Optional +from typing import Any, Dict, List, Optional +from collections.abc import Collection from uuid import uuid4 from psycopg import sql @@ -25,7 +26,7 @@ class GrlIqEventManager: def update_or_create_timing( self, session_uuid: UUIDStr, - timing_data: Optional[TimingData] = None, + timing_data: TimingData | None = None, ) -> PositiveInt: data = { "session_uuid": session_uuid, @@ -77,8 +78,8 @@ class GrlIqEventManager: session_uuid: UUIDStr, event_start: datetime, event_end: datetime, - events: Optional[List[Dict]] = None, - mouse_events: Optional[List[Dict]] = None, + events: list[dict] | None = None, + mouse_events: list[dict] | None = None, ) -> PositiveInt: data = { "uuid": uuid4().hex, @@ -135,14 +136,14 @@ class GrlIqEventManager: def filter( self, - select_str: Optional[str] = None, - session_uuid: Optional[str] = None, - session_uuids: Optional[Collection[str]] = None, - uuids: Optional[Collection[str]] = None, - started_since: Optional[datetime] = None, - limit: Optional[int] = None, + select_str: str | None = None, + session_uuid: str | None = None, + session_uuids: Collection[str] | None = None, + uuids: Collection[str] | None = None, + started_since: datetime | None = None, + limit: int | None = None, order_by: str = "event_start DESC", - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: if not limit: limit = 100 @@ -199,7 +200,7 @@ class GrlIqEventManager: def filter_distinct_timing( self, session_uuids: Collection[str], - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: params = {"session_uuids": list(session_uuids)} query = sql.SQL( """ @@ -229,7 +230,7 @@ class GrlIqEventManager: return res @staticmethod - def process_mouse_events(pointer_moves: List[PointerMove], events: List[Dict]): + def process_mouse_events(pointer_moves: list[PointerMove], events: list[dict]): """ In the db column 'mouse_events' we put all 'pointermove' events. Pull those out, and then any 'pointerdown' and 'pointerup' events from the @@ -274,7 +275,7 @@ class GrlIqEventManager: return mouse_events @staticmethod - def process_keyboard_events(events: List[Dict]): + def process_keyboard_events(events: list[dict]): res = [ KeyboardEvent( type=x["type"], diff --git a/generalresearch/grliq/managers/forensic_results.py b/generalresearch/grliq/managers/forensic_results.py index 30db53d..52bde99 100644 --- a/generalresearch/grliq/managers/forensic_results.py +++ b/generalresearch/grliq/managers/forensic_results.py @@ -1,5 +1,6 @@ from datetime import datetime -from typing import Any, Collection, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple +from collections.abc import Collection from generalresearch.grliq.models.forensic_result import ( GrlIqForensicCategoryResult, @@ -16,16 +17,16 @@ class GrlIqCategoryResultsReader: def filter_category_results( self, - session_uuid: Optional[str] = None, - fingerprint: Optional[str] = None, - phase: Optional[Phase] = None, - uuids: Optional[Collection[str]] = None, - product_ids: Optional[Collection[str]] = None, - created_since: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - limit: Optional[int] = None, - ) -> List[Dict[str, Any]]: + session_uuid: str | None = None, + fingerprint: str | None = None, + phase: Phase | None = None, + uuids: Collection[str] | None = None, + product_ids: Collection[str] | None = None, + created_since: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + limit: int | None = None, + ) -> list[dict[str, Any]]: """ For retrieving GrlIqForensicCategoryResult objects from db. diff --git a/generalresearch/grliq/managers/forensic_summary.py b/generalresearch/grliq/managers/forensic_summary.py index 21b7e4b..5039a38 100644 --- a/generalresearch/grliq/managers/forensic_summary.py +++ b/generalresearch/grliq/managers/forensic_summary.py @@ -2,7 +2,7 @@ from __future__ import annotations import statistics from collections import defaultdict -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from typing import Any, Dict, List import numpy as np @@ -27,7 +27,7 @@ from generalresearch.redis_helper import RedisConfig def calculate_category_summary( - res: List[GrlIqForensicCategoryResult], + res: list[GrlIqForensicCategoryResult], ) -> GrlIqForensicCategorySummary: totals = defaultdict(int) is_complete_count = 0 @@ -55,7 +55,7 @@ def calculate_category_summary( def calculate_checker_summary( - res: List[GrlIqCheckerResults], + res: list[GrlIqCheckerResults], ) -> GrlIqCheckerResultsSummary: totals = defaultdict(list) none_totals = defaultdict(int) @@ -85,8 +85,8 @@ def calculate_checker_summary( def calculate_timing_summary( - redis_config: RedisConfig, timing_res: List[Dict[str, Any]] -) -> Dict[str, TimingDataCountrySummary]: + redis_config: RedisConfig, timing_res: list[dict[str, Any]] +) -> dict[str, TimingDataCountrySummary]: country_median_rtts = defaultdict(list) for x in timing_res: @@ -137,7 +137,7 @@ def run_user_forensic_summary( user: User, ) -> UserForensicSummary: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) created_between = (now - timedelta(days=90), now) select_str = "id, session_uuid, product_id, product_user_id, created_at, result_data, category_result" res = iq_dm.filter( @@ -158,7 +158,7 @@ def run_user_forensic_summary( ) session_uuids = {x["session_uuid"] for x in res} - timing_res: List[Dict] = iq_em.filter_distinct_timing(session_uuids=session_uuids) + timing_res: list[dict] = iq_em.filter_distinct_timing(session_uuids=session_uuids) country_timing_data_summary = ( calculate_timing_summary(redis_config=redis_config, timing_res=timing_res) diff --git a/generalresearch/grliq/models/custom_types.py b/generalresearch/grliq/models/custom_types.py index c5eb93c..5c7f155 100644 --- a/generalresearch/grliq/models/custom_types.py +++ b/generalresearch/grliq/models/custom_types.py @@ -1,5 +1,5 @@ import annotated_types -from typing_extensions import Annotated +from typing import Annotated GrlIqScore = Annotated[int, annotated_types.Ge(0), annotated_types.Le(100)] GrlIqAvgScore = Annotated[float, annotated_types.Ge(0), annotated_types.Le(100)] diff --git a/generalresearch/grliq/models/decider.py b/generalresearch/grliq/models/decider.py index 4464a7f..d24e150 100644 --- a/generalresearch/grliq/models/decider.py +++ b/generalresearch/grliq/models/decider.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from enum import Enum from pydantic import BaseModel, ConfigDict, Field @@ -35,7 +35,7 @@ class GrlIqAttemptResult(BaseModel): timestamp: AwareDatetimeISO = Field( description="When this decision was made", - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), ) decider: Decider = Field(description="Where this decision was made") decision: AttemptDecision = Field( diff --git a/generalresearch/grliq/models/events.py b/generalresearch/grliq/models/events.py index 7e9ebee..5daa032 100644 --- a/generalresearch/grliq/models/events.py +++ b/generalresearch/grliq/models/events.py @@ -14,7 +14,7 @@ from pydantic import ( NonNegativeInt, PositiveFloat, ) -from typing_extensions import Self +from typing import Self from generalresearch.models.custom_types import AwareDatetimeISO, IPvAnyAddressStr @@ -37,14 +37,14 @@ class Event: # in microseconds, since page load (?) timeStamp: float # optional ID of the event target (e.g.: where the mouse is hovering) - _elementId: Optional[str] = None + _elementId: str | None = None # optional tag name of the event target - _elementTagName: Optional[str] = None + _elementTagName: str | None = None # extracted coordinates for the element being interacted with - _elementBounds: Optional[Bounds] = None + _elementBounds: Bounds | None = None @classmethod - def from_dict(cls, data: Dict[str, Any]) -> Self: + def from_dict(cls, data: dict[str, Any]) -> Self: data = {k: v for k, v in data.items() if k in cls.__dataclass_fields__} bounds = data.get("_elementBounds") if bounds is not None and not isinstance(bounds, Bounds): @@ -104,13 +104,13 @@ class KeyboardEvent(Event): # "insertText", "insertCompositionText", "deleteCompositionText", # "insertFromComposition", "deleteContentBackward" - inputType: Optional[str] + inputType: str | None # e.g., 'Enter', 'a', 'Backspace' - key: Optional[str] = None + key: str | None = None # This is the actual text, if applicable - data: Optional[str] = None + data: str | None = None @property def key_text(self): @@ -159,18 +159,18 @@ class TimingData(BaseModel): """ model_config = ConfigDict(extra="forbid", validate_assignment=True) - client_rtts: List[float] = Field() - server_rtts: List[float] = Field() + client_rtts: list[float] = Field() + server_rtts: list[float] = Field() # Have to be optional for backwards-compatibility, but should always be set. - started_at: Optional[AwareDatetimeISO] = Field(default=None) - ended_at: Optional[AwareDatetimeISO] = Field(default=None) - client_ip: Optional[IPvAnyAddressStr] = Field( + started_at: AwareDatetimeISO | None = Field(default=None) + ended_at: AwareDatetimeISO | None = Field(default=None) + client_ip: IPvAnyAddressStr | None = Field( description="This comes from the websocket request's headers", examples=["72.39.217.116"], default=None, ) - server_hostname: Optional[str] = Field( + server_hostname: str | None = Field( description="The hostname of the server that handled this request", examples=["grliq-web-0"], default=None, @@ -189,7 +189,7 @@ class TimingData(BaseModel): def has_data(self): return len(self.client_rtts) > 0 and len(self.server_rtts) > 0 - def filter_rtts(self, rtts: List[float]) -> List[float]: + def filter_rtts(self, rtts: list[float]) -> list[float]: # Skip the first 5 pings, unless we have <10 pings, then get the last # 5 instead. # The first couple pings are usually outliers as they are running @@ -234,7 +234,7 @@ class TimingData(BaseModel): return rtts @property - def summarize(self) -> Optional[TimingDataSummary]: + def summarize(self) -> TimingDataSummary | None: if len(self.filtered_rtts) < 5: return None diff --git a/generalresearch/grliq/models/forensic_data.py b/generalresearch/grliq/models/forensic_data.py index f8bdd98..eda7186 100644 --- a/generalresearch/grliq/models/forensic_data.py +++ b/generalresearch/grliq/models/forensic_data.py @@ -3,7 +3,7 @@ from __future__ import annotations import hashlib import re from collections import Counter -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from enum import Enum from functools import cached_property from typing import Any, Literal @@ -23,7 +23,7 @@ from pydantic import ( ) from pydantic.json_schema import SkipJsonSchema from pydantic_extra_types.timezone_name import TimeZoneName -from typing_extensions import Annotated, Self +from typing import Annotated, Self from generalresearch.grliq.models import ( AUDIO_CODEC_NAMES, @@ -776,7 +776,7 @@ class GrlIqData(BaseModel): ), "product_user_id mismatch" # validate the Session's mid is "recent" - assert (datetime.now(tz=timezone.utc) - session.started) < timedelta( + assert (datetime.now(tz=UTC) - session.started) < timedelta( minutes=90 ), "expired session" diff --git a/generalresearch/grliq/models/forensic_summary.py b/generalresearch/grliq/models/forensic_summary.py index d6f46f8..5ecf1a4 100644 --- a/generalresearch/grliq/models/forensic_summary.py +++ b/generalresearch/grliq/models/forensic_summary.py @@ -228,7 +228,7 @@ class CountryRTTDistribution(BaseModel): rtt_mean: float = Field(gt=0, examples=[179.302]) rtt_max: float = Field(gt=0, examples=[890.006]) rtt_std: float = Field(gt=0, examples=[46.831]) - rtt_percentiles: List[float] = Field( + rtt_percentiles: list[float] = Field( min_length=101, max_length=101, examples=[example_rtt_percentiles] ) diff --git a/generalresearch/grliq/models/useragents.py b/generalresearch/grliq/models/useragents.py index 1953f6d..4bb340e 100644 --- a/generalresearch/grliq/models/useragents.py +++ b/generalresearch/grliq/models/useragents.py @@ -4,7 +4,7 @@ import hashlib from enum import Enum from pydantic import BaseModel, ConfigDict, Field, field_validator -from typing_extensions import Self +from typing import Self from user_agents import parse as ua_parse from user_agents.parsers import UserAgent diff --git a/generalresearch/grliq/utils.py b/generalresearch/grliq/utils.py index 95390a8..ca8c6a1 100644 --- a/generalresearch/grliq/utils.py +++ b/generalresearch/grliq/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from pathlib import Path from uuid import UUID @@ -16,7 +16,7 @@ def get_screenshot_fp( grliq_ss_dir_name: str = "canvas2html", create_dir_if_not_exists: bool = True, ) -> Path | None: - assert created_at.tzinfo == timezone.utc + assert created_at.tzinfo == UTC if isinstance(forensic_uuid, UUID): forensic_uuid = forensic_uuid.hex diff --git a/generalresearch/grpc.py b/generalresearch/grpc.py index 040fd26..178521e 100644 --- a/generalresearch/grpc.py +++ b/generalresearch/grpc.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from google.protobuf.duration_pb2 import Duration from google.protobuf.timestamp_pb2 import Timestamp @@ -20,13 +20,13 @@ def timestamp_from_datetime_nullable(dt: datetime | None) -> Timestamp: def timestamp_to_datetime(ts: Timestamp) -> datetime: - return datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=timezone.utc) + return datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=UTC) def timestamp_to_datetime_nullable(ts: Timestamp) -> datetime | None: # grpc has no None. If a google.protobuf.Timestamp field is not set, it gets interpreted as timestamp 0 - default = datetime.fromtimestamp(0, tz=timezone.utc) - d = datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=timezone.utc) + default = datetime.fromtimestamp(0, tz=UTC) + d = datetime.fromtimestamp(ts.seconds + ts.nanos / 1e9, tz=UTC) return None if d == default else d diff --git a/generalresearch/incite/base.py b/generalresearch/incite/base.py index a8088ac..6888ca2 100644 --- a/generalresearch/incite/base.py +++ b/generalresearch/incite/base.py @@ -8,7 +8,7 @@ import shutil import subprocess import warnings from concurrent.futures import Future -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from os import R_OK, access, listdir from os.path import isdir from os.path import join as pjoin @@ -17,9 +17,8 @@ from sys import platform from typing import ( TYPE_CHECKING, Any, - Callable, - Sequence, ) +from collections.abc import Callable, Sequence from uuid import uuid4 import dask @@ -43,7 +42,7 @@ from pydantic import ( ) from pydantic.json_schema import SkipJsonSchema from sentry_sdk import capture_exception -from typing_extensions import Self +from typing import Self from generalresearch.config import is_debug from generalresearch.incite.schemas import ( @@ -166,7 +165,7 @@ class CollectionBase(BaseModel): offset: str = Field(default="72h", max_length=5) start: AwareDatetimeISO = Field( - default=datetime(year=2018, month=1, day=1, tzinfo=timezone.utc), + default=datetime(year=2018, month=1, day=1, tzinfo=UTC), description="This is the starting point in which data will be retrieved" "in chunks from.", frozen=True, @@ -208,7 +207,7 @@ class CollectionBase(BaseModel): return self offset_total_sec = pd.Timedelta(self.offset).total_seconds() - start_total_sec = (datetime.now(tz=timezone.utc) - self.start).total_seconds() + start_total_sec = (datetime.now(tz=UTC) - self.start).total_seconds() if offset_total_sec > start_total_sec: raise ValueError("Offset must be equal to, or smaller the start timestamp") @@ -294,14 +293,14 @@ class CollectionBase(BaseModel): @property def interval_range(self) -> list[tuple[datetime, datetime]]: """closed='left', so 0 <= x < 5""" - end = self.finished or datetime.now(tz=timezone.utc).replace(microsecond=0) + end = self.finished or datetime.now(tz=UTC).replace(microsecond=0) iv_r = self._interval_range(end) return [(iv.left.to_pydatetime(), iv.right.to_pydatetime()) for iv in iv_r] @property def progress(self) -> pd.DataFrame: records = [i.to_dict() for i in self.items] - end = self.finished if self.finished else datetime.now(tz=timezone.utc) + end = self.finished if self.finished else datetime.now(tz=UTC) return pd.DataFrame.from_records(records, index=self._interval_range(end)) @property @@ -626,15 +625,15 @@ class CollectionBase(BaseModel): return res def get_items_from_year(self, year: int) -> Items: - ts = datetime(year=year, month=1, day=1, tzinfo=timezone.utc) + ts = datetime(year=year, month=1, day=1, tzinfo=UTC) return self.get_items(since=ts) def get_items_last90(self) -> Items: - ts = datetime.now(tz=timezone.utc) - timedelta(days=90) + ts = datetime.now(tz=UTC) - timedelta(days=90) return self.get_items(since=ts) def get_items_last365(self) -> Items: - ts = datetime.now(tz=timezone.utc) - timedelta(days=365) + ts = datetime.now(tz=UTC) - timedelta(days=365) return self.get_items(since=ts) @@ -642,7 +641,7 @@ class CollectionItemBase(BaseModel): # I want to intentionally keep these as native python types, and not # pandas specific types. start: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc).replace(microsecond=0) + default_factory=lambda: datetime.now(tz=UTC).replace(microsecond=0) ) # --- Private attrs --- @@ -839,7 +838,7 @@ class CollectionItemBase(BaseModel): if archive_after is None: return False - return datetime.now(tz=timezone.utc) > self.finish + archive_after + return datetime.now(tz=UTC) > self.finish + archive_after def set_empty(self): assert ( diff --git a/generalresearch/incite/defaults.py b/generalresearch/incite/defaults.py index 421710e..5a95607 100644 --- a/generalresearch/incite/defaults.py +++ b/generalresearch/incite/defaults.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from generalresearch.incite.base import GRLDatasets from generalresearch.incite.collections import DFCollectionType @@ -37,69 +37,69 @@ from generalresearch.sql_helper import SqlHelper def session_df_collection( - ds: "GRLDatasets", pg_config: PostgresConfig + ds: GRLDatasets, pg_config: PostgresConfig ) -> SessionDFCollection: return SessionDFCollection( offset="37h", pg_config=pg_config, - start=datetime(year=2022, month=5, day=3, hour=12, tzinfo=timezone.utc), + start=datetime(year=2022, month=5, day=3, hour=12, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.SESSION), ) def wall_df_collection( - ds: "GRLDatasets", pg_config: PostgresConfig + ds: GRLDatasets, pg_config: PostgresConfig ) -> WallDFCollection: return WallDFCollection( offset="49h", pg_config=pg_config, - start=datetime(year=2022, month=5, day=3, hour=12, tzinfo=timezone.utc), + start=datetime(year=2022, month=5, day=3, hour=12, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.WALL), ) def user_df_collection( - ds: "GRLDatasets", pg_config: PostgresConfig + ds: GRLDatasets, pg_config: PostgresConfig ) -> UserDFCollection: return UserDFCollection( offset="73h", pg_config=pg_config, - start=datetime(year=2016, month=7, day=13, hour=1, tzinfo=timezone.utc), + start=datetime(year=2016, month=7, day=13, hour=1, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.USER), ) def task_df_collection( - ds: "GRLDatasets", pg_config: PostgresConfig + ds: GRLDatasets, pg_config: PostgresConfig ) -> TaskAdjustmentDFCollection: return TaskAdjustmentDFCollection( offset="48h", pg_config=pg_config, - start=datetime(year=2022, month=7, day=16, hour=0, tzinfo=timezone.utc), + start=datetime(year=2022, month=7, day=16, hour=0, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.TASK_ADJUSTMENT), ) def ledger_df_collection( - ds: "GRLDatasets", pg_config: PostgresConfig + ds: GRLDatasets, pg_config: PostgresConfig ) -> LedgerDFCollection: return LedgerDFCollection( offset="12d", pg_config=pg_config, # thl_web:ledger_transaction - 1st record is 2018-03-14 20:22:17.408232 - start=datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc), + start=datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.LEDGER), ) # --- Marketplace Specifics --- # def innovate_survey_history_collection( - ds: "GRLDatasets", sql_helper: SqlHelper + ds: GRLDatasets, sql_helper: SqlHelper ) -> InnovateSurveyHistoryCollection: return InnovateSurveyHistoryCollection( offset="12h", sql_helper=sql_helper, - start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=timezone.utc), + start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=UTC), archive_path=ds.archive_path( enum_type=DFCollectionType.INNOVATE_SURVEY_HISTORY ), @@ -107,12 +107,12 @@ def innovate_survey_history_collection( def morning_survey_ts_collection( - ds: "GRLDatasets", sql_helper: SqlHelper + ds: GRLDatasets, sql_helper: SqlHelper ) -> MorningSurveyTimeseriesCollection: return MorningSurveyTimeseriesCollection( offset="12h", sql_helper=sql_helper, - start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=timezone.utc), + start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=UTC), archive_path=ds.archive_path( enum_type=DFCollectionType.MORNING_SURVEY_TIMESERIES ), @@ -120,23 +120,23 @@ def morning_survey_ts_collection( def sago_survey_history_collection( - ds: "GRLDatasets", sql_helper: SqlHelper + ds: GRLDatasets, sql_helper: SqlHelper ) -> SagoSurveyHistoryCollection: return SagoSurveyHistoryCollection( offset="12h", sql_helper=sql_helper, - start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=timezone.utc), + start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=UTC), archive_path=ds.archive_path(enum_type=DFCollectionType.SAGO_SURVEY_HISTORY), ) def spectrum_survey_ts_collection( - ds: "GRLDatasets", sql_helper: SqlHelper + ds: GRLDatasets, sql_helper: SqlHelper ) -> SpectrumSurveyTimeseriesCollection: return SpectrumSurveyTimeseriesCollection( offset="12h", sql_helper=sql_helper, - start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=timezone.utc), + start=datetime(year=2024, month=3, day=1, hour=0, tzinfo=UTC), archive_path=ds.archive_path( enum_type=DFCollectionType.SPECTRUM_SURVEY_TIMESERIES ), @@ -144,50 +144,50 @@ def spectrum_survey_ts_collection( # --- Mergers: Foundations --- # -def user_id_product(ds: "GRLDatasets") -> UserIdProductMerge: +def user_id_product(ds: GRLDatasets) -> UserIdProductMerge: return UserIdProductMerge( - start=datetime(year=2010, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=2010, month=1, day=1, tzinfo=UTC), offset=None, archive_path=ds.archive_path(enum_type=MergeType.USER_ID_PRODUCT), ) -def enriched_session(ds: "GRLDatasets") -> EnrichedSessionMerge: +def enriched_session(ds: GRLDatasets) -> EnrichedSessionMerge: return EnrichedSessionMerge( - start=datetime(year=2023, month=5, day=1, tzinfo=timezone.utc), + start=datetime(year=2023, month=5, day=1, tzinfo=UTC), offset="14d", archive_path=ds.archive_path(enum_type=MergeType.ENRICHED_SESSION), ) -def enriched_wall(ds: "GRLDatasets") -> EnrichedWallMerge: +def enriched_wall(ds: GRLDatasets) -> EnrichedWallMerge: return EnrichedWallMerge( # start=datetime(year=2022, month=5, day=1, tzinfo=timezone.utc), - start=datetime(year=2023, month=7, day=23, tzinfo=timezone.utc), + start=datetime(year=2023, month=7, day=23, tzinfo=UTC), offset="14d", archive_path=ds.archive_path(enum_type=MergeType.ENRICHED_WALL), ) -def enriched_task_adjust(ds: "GRLDatasets") -> EnrichedTaskAdjustMerge: +def enriched_task_adjust(ds: GRLDatasets) -> EnrichedTaskAdjustMerge: return EnrichedTaskAdjustMerge( - start=datetime(year=2010, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=2010, month=1, day=1, tzinfo=UTC), offset=None, archive_path=ds.archive_path(enum_type=MergeType.ENRICHED_TASK_ADJUST), ) # --- Mergers: Others --- # -def pop_ledger(ds: "GRLDatasets") -> PopLedgerMerge: +def pop_ledger(ds: GRLDatasets) -> PopLedgerMerge: return PopLedgerMerge( # thl_web:ledger_transaction - 1st record is 2018-03-14 20:22:17.408232 - start=datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc), + start=datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC), offset="30d", archive_path=ds.archive_path(enum_type=MergeType.POP_LEDGER), ) -def ym_survey_wall(ds: "GRLDatasets") -> YMSurveyWallMerge: +def ym_survey_wall(ds: GRLDatasets) -> YMSurveyWallMerge: return YMSurveyWallMerge( start=None, offset="10D", diff --git a/generalresearch/incite/mergers/__init__.py b/generalresearch/incite/mergers/__init__.py index b9c3789..9c5276c 100644 --- a/generalresearch/incite/mergers/__init__.py +++ b/generalresearch/incite/mergers/__init__.py @@ -1,7 +1,7 @@ import logging import os.path import subprocess -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from enum import Enum from sys import platform from typing import List, Optional, Type @@ -11,7 +11,7 @@ import pandas as pd from dask.distributed import Client from pandera.pandas import DataFrameSchema from pydantic import Field, ValidationInfo, field_validator, model_validator -from typing_extensions import Self +from typing import Self from generalresearch.incite.base import CollectionBase, CollectionItemBase from generalresearch.incite.schemas import PARTITION_ON @@ -80,7 +80,7 @@ class MergeCollectionItem(CollectionItemBase): pd.Timestamp(self.start) + pd.Timedelta(self._collection.offset) ).to_pydatetime() else: - return datetime.now(tz=timezone.utc).replace(microsecond=0) + return datetime.now(tz=UTC).replace(microsecond=0) @property def filename(self) -> str: @@ -236,20 +236,20 @@ class MergeCollection(CollectionBase): # In a merge, we can set offset = None which indicates that there is only 1 # period/item where the range is 'start' until now. - offset: Optional[str] = Field(default="72h") + offset: str | None = Field(default="72h") # In a merge, we can set start = None which indicates that there is only 1 # period/item where the range is (now - offset) until now. - start: Optional[AwareDatetimeISO] = Field( + start: AwareDatetimeISO | None = Field( default=None, description="This is the starting point in which data will" " be retrieved in chunks from.", frozen=True, ) - merge_type: Optional[MergeType] = Field(default=None) - group_by: Optional[str] = Field(default=None) - grouped_key: Optional[str] = Field(default=None) - collection_item_class: Type[MergeCollectionItem] = MergeCollectionItem + merge_type: MergeType | None = Field(default=None) + group_by: str | None = Field(default=None) + grouped_key: str | None = Field(default=None) + collection_item_class: type[MergeCollectionItem] = MergeCollectionItem @model_validator(mode="after") def check_start_and_offset_nullable(self) -> Self: @@ -269,16 +269,16 @@ class MergeCollection(CollectionBase): # --- Properties --- @property - def interval_start(self) -> Optional[datetime]: + def interval_start(self) -> datetime | None: # if self.start is None and self.offset is set, the inferred start is (now - offset) if self.start is None: - return datetime.now(tz=timezone.utc).replace(microsecond=0) - pd.Timedelta( + return datetime.now(tz=UTC).replace(microsecond=0) - pd.Timedelta( self.offset ) return self.start @property - def items(self) -> List[MergeCollectionItem]: + def items(self) -> list[MergeCollectionItem]: items = [] for iv in self.interval_range: cm = self.collection_item_class(start=iv[0]) diff --git a/generalresearch/incite/mergers/ym_wall_summary.py b/generalresearch/incite/mergers/ym_wall_summary.py index 2f5995f..4816c05 100644 --- a/generalresearch/incite/mergers/ym_wall_summary.py +++ b/generalresearch/incite/mergers/ym_wall_summary.py @@ -82,7 +82,7 @@ class YMWallSummaryMergeItem(MergeCollectionItem): class YMWallSummaryMerge(MergeCollection): merge_type: Literal[MergeType.YM_WALL_SUMMARY] = MergeType.YM_WALL_SUMMARY _schema = YMWallSummarySchema - collection_item_class: Type[YMWallSummaryMergeItem] = YMWallSummaryMergeItem + collection_item_class: type[YMWallSummaryMergeItem] = YMWallSummaryMergeItem items: list[YMWallSummaryMergeItem] = Field(default_factory=list) @field_validator("offset") diff --git a/generalresearch/incite/schemas/__init__.py b/generalresearch/incite/schemas/__init__.py index c0000d1..6fc83b0 100644 --- a/generalresearch/incite/schemas/__init__.py +++ b/generalresearch/incite/schemas/__init__.py @@ -11,8 +11,8 @@ ARCHIVE_AFTER = "archive_after" PARTITION_ON = "partition_on" -def empty_dataframe_from_schema(schema: pa.DataFrameSchema) -> "pd.DataFrame": - index_names: List[str] = schema.index.names +def empty_dataframe_from_schema(schema: pa.DataFrameSchema) -> pd.DataFrame: + index_names: list[str] = schema.index.names columns = set(schema.dtypes.keys()) if len(index_names) > 1: diff --git a/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py b/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py index cc909f6..97e73a3 100644 --- a/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py +++ b/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py @@ -13,7 +13,7 @@ from generalresearch.models.thl.definitions import ( thl_task_adj_columns = THLTaskAdjustmentSchema.columns.copy() -COUNTRY_ISOS: Set[str] = Localelator().get_all_countries() +COUNTRY_ISOS: set[str] = Localelator().get_all_countries() kosovo = "xk" COUNTRY_ISOS.add(kosovo) BIGINT = 9223372036854775807 diff --git a/generalresearch/incite/schemas/thl_web.py b/generalresearch/incite/schemas/thl_web.py index 5073a18..b831b9a 100644 --- a/generalresearch/incite/schemas/thl_web.py +++ b/generalresearch/incite/schemas/thl_web.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index, MultiIndex @@ -105,7 +105,7 @@ THLWallSchema = DataFrameSchema( ), "started": Column( dtype=pd.DatetimeTZDtype(tz="UTC"), - checks=[Check(lambda x: x < datetime.now(tz=timezone.utc))], + checks=[Check(lambda x: x < datetime.now(tz=UTC))], nullable=False, ), "session_id": Column( @@ -205,12 +205,12 @@ THLSessionSchema = DataFrameSchema( ), "started": Column( dtype=pd.DatetimeTZDtype(tz="UTC"), - checks=[Check(lambda x: x < datetime.now(tz=timezone.utc))], + checks=[Check(lambda x: x < datetime.now(tz=UTC))], nullable=True, ), "finished": Column( dtype=pd.DatetimeTZDtype(tz="UTC"), - checks=[Check(lambda x: x < datetime.now(tz=timezone.utc))], + checks=[Check(lambda x: x < datetime.now(tz=UTC))], nullable=True, ), "loi_min": Column(dtype="Int64", nullable=True), @@ -450,7 +450,7 @@ THLTaskAdjustmentSchema = DataFrameSchema( ), "started": Column( dtype=pd.DatetimeTZDtype(tz="UTC"), - checks=[Check(lambda x: x < datetime.now(tz=timezone.utc))], + checks=[Check(lambda x: x < datetime.now(tz=UTC))], ), "source": Column( dtype=str, diff --git a/generalresearch/locales/__init__.py b/generalresearch/locales/__init__.py index 88b72e6..813966e 100644 --- a/generalresearch/locales/__init__.py +++ b/generalresearch/locales/__init__.py @@ -43,11 +43,11 @@ class Localelator: pkgutil.get_data(__name__, "country_default_lang.json") ) - def get_all_languages(self) -> Set[str]: + def get_all_languages(self) -> set[str]: # returns only the ISO 639-2/B (three-letter codes) return set(self.lang_alpha2_to_alpha3b.values()) - def get_all_countries(self) -> Set[str]: + def get_all_countries(self) -> set[str]: # returns only the ISO 3166-1 alpha-2 (two-letter codes) return set(self.country_alpha3_to_alpha2.values()) diff --git a/generalresearch/locales/timezone.py b/generalresearch/locales/timezone.py index 50d539d..fce6e0e 100644 --- a/generalresearch/locales/timezone.py +++ b/generalresearch/locales/timezone.py @@ -3,7 +3,7 @@ from typing import Optional from pytz import country_timezones -def get_default_timezone(country_iso: str) -> Optional[str]: +def get_default_timezone(country_iso: str) -> str | None: # to list all: # from pytz import country_names, country_timezones # [country_timezones.get(country) for country in country_names] @@ -72,6 +72,6 @@ country_default_locale = { } -def get_default_locale(country_iso: str) -> Optional[str]: +def get_default_locale(country_iso: str) -> str | None: # todo: "https://cdn.simplelocalize.io/public/v1/locales" to fill in the rest? return country_default_locale.get(country_iso, None) diff --git a/generalresearch/managers/cint/profiling.py b/generalresearch/managers/cint/profiling.py index d549e94..9216aa5 100644 --- a/generalresearch/managers/cint/profiling.py +++ b/generalresearch/managers/cint/profiling.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Collection +from collections.abc import Collection from generalresearch.models.cint.question import CintQuestion from generalresearch.sql_helper import SqlHelper diff --git a/generalresearch/managers/cint/survey.py b/generalresearch/managers/cint/survey.py index f80542e..da1ecd9 100644 --- a/generalresearch/managers/cint/survey.py +++ b/generalresearch/managers/cint/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -107,7 +107,7 @@ class CintSurveyManager(SurveyManager): return True def update(self, surveys: list[CintSurvey]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) for survey in surveys: survey.last_updated = now diff --git a/generalresearch/managers/criteria.py b/generalresearch/managers/criteria.py index fe70732..b5d9830 100644 --- a/generalresearch/managers/criteria.py +++ b/generalresearch/managers/criteria.py @@ -2,7 +2,7 @@ from __future__ import annotations from abc import ABC from collections.abc import Collection -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from more_itertools import chunked @@ -65,7 +65,7 @@ class CriteriaManager(SqlManager, ABC): new_hashes = this_hashes - known_hashes if new_hashes: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) values = [ condition.to_mysql() for condition in conditions diff --git a/generalresearch/managers/dynata/profiling.py b/generalresearch/managers/dynata/profiling.py index 25afbe2..661bdad 100644 --- a/generalresearch/managers/dynata/profiling.py +++ b/generalresearch/managers/dynata/profiling.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Collection +from collections.abc import Collection from generalresearch.models.dynata.question import DynataQuestion from generalresearch.sql_helper import SqlHelper diff --git a/generalresearch/managers/dynata/survey.py b/generalresearch/managers/dynata/survey.py index 372a57d..3a15c4d 100644 --- a/generalresearch/managers/dynata/survey.py +++ b/generalresearch/managers/dynata/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -102,7 +102,7 @@ class DynataSurveyManager(SurveyManager): return surveys def create(self, survey: DynataSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -123,7 +123,7 @@ class DynataSurveyManager(SurveyManager): return True def update(self, surveys: list[DynataSurvey]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) update_fields = self.SURVEY_FIELDS + ["last_updated"] data = [survey.to_mysql() for survey in surveys] diff --git a/generalresearch/managers/events.py b/generalresearch/managers/events.py index 0be2bb9..f3c6a04 100644 --- a/generalresearch/managers/events.py +++ b/generalresearch/managers/events.py @@ -5,7 +5,7 @@ import math import socket import threading import time -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from typing import TYPE_CHECKING @@ -141,7 +141,7 @@ class UserStatsManager(RedisManager): pipe.execute() def mark_user_active(self, user: User) -> None: - now = datetime.now(tz=timezone.utc).isoformat() + now = datetime.now(tz=UTC).isoformat() r = self.redis_client pipe = r.pipeline(transaction=False) @@ -175,7 +175,7 @@ class UserStatsManager(RedisManager): # This call is idempotent; it can be called multiple times (for the # same user) and won't falsely increase a counter; it will just # reset the expiration for this user (times out after 60 min) - now = datetime.now(tz=timezone.utc).isoformat() + now = datetime.now(tz=UTC).isoformat() r = self.redis_client pipe = r.pipeline(transaction=False) diff --git a/generalresearch/managers/gr/authentication.py b/generalresearch/managers/gr/authentication.py index a402693..409cb10 100644 --- a/generalresearch/managers/gr/authentication.py +++ b/generalresearch/managers/gr/authentication.py @@ -3,7 +3,7 @@ from __future__ import annotations import binascii import logging import os -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from typing import TYPE_CHECKING, Any from psycopg import sql @@ -29,7 +29,7 @@ class GRUserManager(PostgresManagerWithRedis): ) -> GRUser: from generalresearch.models.gr.authentication import GRUser - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) instance = GRUser.model_validate( { @@ -147,7 +147,7 @@ class GRUserManager(PostgresManagerWithRedis): for item in res: for k, v in item.items(): if isinstance(item[k], datetime): - item[k] = item[k].replace(tzinfo=timezone.utc) + item[k] = item[k].replace(tzinfo=UTC) return [GRUser.model_validate(item) for item in res] @@ -216,7 +216,7 @@ class GRTokenManager(PostgresManager): "key": api_key, "user_id": gr_user.id, "user": gr_user, - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), } ) @@ -251,7 +251,7 @@ class GRTokenManager(PostgresManager): token = GRToken.model_validate( { "key": binascii.hexlify(os.urandom(20)).decode(), - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), "user_id": user_id, } ) @@ -298,6 +298,6 @@ class GRTokenManager(PostgresManager): for k, _ in res.items(): if isinstance(res[k], datetime): - res[k] = res[k].replace(tzinfo=timezone.utc) + res[k] = res[k].replace(tzinfo=UTC) return GRToken.model_validate(res) diff --git a/generalresearch/managers/gr/team.py b/generalresearch/managers/gr/team.py index 6de82b0..393f446 100644 --- a/generalresearch/managers/gr/team.py +++ b/generalresearch/managers/gr/team.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from typing import TYPE_CHECKING from uuid import uuid4 @@ -43,7 +43,7 @@ class MembershipManager(PostgresManager): owner=False, team_id=team.id, user_id=gr_user.id, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) data = membership.model_dump(by_alias=True) diff --git a/generalresearch/managers/innovate/survey.py b/generalresearch/managers/innovate/survey.py index 7db2f49..c65b100 100644 --- a/generalresearch/managers/innovate/survey.py +++ b/generalresearch/managers/innovate/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -121,7 +121,7 @@ class InnovateSurveyManager(SurveyManager): return surveys def create(self, survey: InnovateSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -142,7 +142,7 @@ class InnovateSurveyManager(SurveyManager): return True def update(self, surveys: list[InnovateSurvey]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) update_fields = self.SURVEY_FIELDS + ["updated"] data = [survey.to_mysql() for survey in surveys] diff --git a/generalresearch/managers/leaderboard/manager.py b/generalresearch/managers/leaderboard/manager.py index 0bf0312..71c3a73 100644 --- a/generalresearch/managers/leaderboard/manager.py +++ b/generalresearch/managers/leaderboard/manager.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from decimal import Decimal from functools import cached_property from typing import TYPE_CHECKING @@ -45,7 +45,7 @@ class LeaderboardManager: self.country_iso = country_iso self.within_time_aware = None if within_time is None: - self.within_time_aware = datetime.now(tz=timezone.utc).astimezone( + self.within_time_aware = datetime.now(tz=UTC).astimezone( self.timezone ) elif within_time.tzinfo is not None: @@ -57,7 +57,7 @@ class LeaderboardManager: @cached_property def period(self) -> Period: local_ts = self.within_time_aware - assert local_ts.tzinfo != timezone.utc and local_ts.tzinfo is not None + assert local_ts.tzinfo != UTC and local_ts.tzinfo is not None t = pd.Timestamp(local_ts).tz_localize(tz=None) freq_pd = { LeaderboardFrequency.WEEKLY: "W-SUN", diff --git a/generalresearch/managers/morning/survey.py b/generalresearch/managers/morning/survey.py index 2d86f0f..5fba70d 100644 --- a/generalresearch/managers/morning/survey.py +++ b/generalresearch/managers/morning/survey.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -138,7 +138,7 @@ class MorningSurveyManager(SurveyManager): return bids def create(self, bid: MorningBid) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = bid.to_mysql() create_fields = self.BID_FIELDS + ["created", "updated"] @@ -179,14 +179,14 @@ class MorningSurveyManager(SurveyManager): return True def update(self, surveys: list[MorningBid]) -> None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) for survey in surveys: self.update_one(survey, now=now) def update_one(self, bid: MorningBid, now: datetime | None = None) -> bool: if now is None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = bid.to_mysql() d["updated"] = now diff --git a/generalresearch/managers/network/label.py b/generalresearch/managers/network/label.py index 1efe875..c5306a6 100644 --- a/generalresearch/managers/network/label.py +++ b/generalresearch/managers/network/label.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from psycopg import sql from pydantic import IPvAnyNetwork, TypeAdapter @@ -48,8 +48,8 @@ class IPLabelManager(PostgresManager): filters = [] params = {} if labeled_after or labeled_before: - time_end = labeled_before or datetime.now(tz=timezone.utc) - time_start = labeled_after or datetime(2017, 1, 1, tzinfo=timezone.utc) + time_end = labeled_before or datetime.now(tz=UTC) + time_start = labeled_after or datetime(2017, 1, 1, tzinfo=UTC) assert time_start.tzinfo.utcoffset(time_start) == timedelta(), "must be UTC" assert time_end.tzinfo.utcoffset(time_end) == timedelta(), "must be UTC" filters.append("labeled_at BETWEEN %(time_start)s AND %(time_end)s") diff --git a/generalresearch/managers/precision/survey.py b/generalresearch/managers/precision/survey.py index 6fb30f2..833cb28 100644 --- a/generalresearch/managers/precision/survey.py +++ b/generalresearch/managers/precision/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -104,7 +104,7 @@ class PrecisionSurveyManager(SurveyManager): return surveys def create(self, survey: PrecisionSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(False) @@ -151,7 +151,7 @@ class PrecisionSurveyManager(SurveyManager): return True def update_one(self, survey: PrecisionSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() d["updated"] = now diff --git a/generalresearch/managers/prodege/survey.py b/generalresearch/managers/prodege/survey.py index f555290..750383f 100644 --- a/generalresearch/managers/prodege/survey.py +++ b/generalresearch/managers/prodege/survey.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql @@ -93,7 +93,7 @@ class ProdegeSurveyManager(SurveyManager): return surveys def create(self, survey: ProdegeSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -114,7 +114,7 @@ class ProdegeSurveyManager(SurveyManager): return True def update(self, surveys: list[ProdegeSurvey]) -> None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) # Do to stupidity with bid/actual loi/ir values (see ProdegeSurvey.to_mysql), we now # can't do a bulk update b/c the fields may be different in different rows. Just do @@ -124,7 +124,7 @@ class ProdegeSurveyManager(SurveyManager): def update_one(self, survey: ProdegeSurvey, now=None) -> bool: if now is None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() # We have to have special logic for bid/actual loi/ir here. The api is # stupid and only returns one set of them. If we just do the db diff --git a/generalresearch/managers/repdata/survey.py b/generalresearch/managers/repdata/survey.py index 2e2224f..1e1c3c6 100644 --- a/generalresearch/managers/repdata/survey.py +++ b/generalresearch/managers/repdata/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql @@ -122,7 +122,7 @@ class RepDataSurveyManager(SurveyManager): return list(surveys.values()) def create(self, survey: RepDataSurvey | RepDataSurveyHashed) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -160,7 +160,7 @@ class RepDataSurveyManager(SurveyManager): return True def update(self, surveys: list[RepDataSurveyHashed]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) update_fields = self.SURVEY_FIELDS + ["last_updated"] data = [survey.to_mysql() for survey in surveys] diff --git a/generalresearch/managers/sago/survey.py b/generalresearch/managers/sago/survey.py index 325639f..2582902 100644 --- a/generalresearch/managers/sago/survey.py +++ b/generalresearch/managers/sago/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -101,7 +101,7 @@ class SagoSurveyManager(SurveyManager): return surveys def create(self, survey: SagoSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -122,7 +122,7 @@ class SagoSurveyManager(SurveyManager): return True def update(self, surveys: list[SagoSurvey]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) update_fields = self.SURVEY_FIELDS + ["updated"] data = [survey.to_mysql() for survey in surveys] @@ -131,7 +131,7 @@ class SagoSurveyManager(SurveyManager): return True def update_field(self, survey: SagoSurvey, field: str) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) conn: pymysql.Connection = self.sql_helper.make_connection() value = survey.to_mysql()[field] c = conn.cursor() diff --git a/generalresearch/managers/spectrum/survey.py b/generalresearch/managers/spectrum/survey.py index 3ff2db8..58f8a1a 100644 --- a/generalresearch/managers/spectrum/survey.py +++ b/generalresearch/managers/spectrum/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC import pymysql from pymysql import IntegrityError @@ -110,7 +110,7 @@ class SpectrumSurveyManager(SurveyManager): return surveys def create(self, survey: SpectrumSurvey) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() conn: pymysql.Connection = self.sql_helper.make_connection() conn.autocommit(True) @@ -134,7 +134,7 @@ class SpectrumSurveyManager(SurveyManager): return True def update(self, surveys: list[SpectrumSurvey]) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) # Due to stupidity with bid/actual loi/ir values (last block nonsense), # we can't do a bulk update b/c the fields may be different in @@ -146,7 +146,7 @@ class SpectrumSurveyManager(SurveyManager): def update_one(self, survey: SpectrumSurvey, now: datetime | None = None) -> bool: if now is None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = survey.to_mysql() # We have to have special logic for bid/actual loi/ir here. The api diff --git a/generalresearch/managers/thl/buyer.py b/generalresearch/managers/thl/buyer.py index 04452cd..ae40bc8 100644 --- a/generalresearch/managers/thl/buyer.py +++ b/generalresearch/managers/thl/buyer.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from generalresearch.managers.base import Permission, PostgresManager from generalresearch.models import Source @@ -45,7 +45,7 @@ class BuyerManager(PostgresManager): return None def bulk_get_or_create(self, source: Source, codes: Collection[str]) -> list[Buyer]: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) buyers = [] params_seq = [] diff --git a/generalresearch/managers/thl/cashout_method.py b/generalresearch/managers/thl/cashout_method.py index 7365878..94617f7 100644 --- a/generalresearch/managers/thl/cashout_method.py +++ b/generalresearch/managers/thl/cashout_method.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Collection from copy import copy -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from typing import Any from uuid import UUID, uuid4 @@ -21,7 +21,7 @@ from generalresearch.models.thl.wallet.cashout_method import ( class CashoutMethodManager(PostgresManager): def create(self, cm: CashoutMethod) -> None: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) query = """ INSERT INTO accounting_cashoutmethod ( id, last_updated, is_live, provider, diff --git a/generalresearch/managers/thl/contest_manager.py b/generalresearch/managers/thl/contest_manager.py index 517f677..286de3d 100644 --- a/generalresearch/managers/thl/contest_manager.py +++ b/generalresearch/managers/thl/contest_manager.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from typing import Any, Literal, cast from uuid import UUID @@ -199,10 +199,10 @@ class ContestBaseManager(PostgresManager): params["contest_type"] = contest_type.value filters.append("contest_type = %(contest_type)s") if starts_at_before is True: - params["starts_at"] = datetime.now(tz=timezone.utc) + params["starts_at"] = datetime.now(tz=UTC) filters.append("starts_at < %(starts_at)s") elif starts_at_before: - assert starts_at_before.tzinfo == timezone.utc + assert starts_at_before.tzinfo == UTC params["starts_at"] = starts_at_before filters.append("starts_at < %(starts_at)s") if name is not None: @@ -822,7 +822,7 @@ class MilestoneContestManager(ContestBaseManager): if decision: contest.update( status=ContestStatus.COMPLETED, - ended_at=datetime.now(tz=timezone.utc), + ended_at=datetime.now(tz=UTC), end_reason=reason, ) self.end_milestone_contest(contest) diff --git a/generalresearch/managers/thl/ledger_manager/conditions.py b/generalresearch/managers/thl/ledger_manager/conditions.py index a457f30..f79a0f0 100644 --- a/generalresearch/managers/thl/ledger_manager/conditions.py +++ b/generalresearch/managers/thl/ledger_manager/conditions.py @@ -1,8 +1,9 @@ from __future__ import annotations import logging -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Callable +from datetime import datetime, timedelta, timezone, UTC +from typing import TYPE_CHECKING +from collections.abc import Callable from generalresearch.config import JAMES_BILLINGS_BPID, JAMES_BILLINGS_TX_CUTOFF from generalresearch.currency import USDCent @@ -73,7 +74,7 @@ def generate_condition_bp_payout( skip_one_per_day_check: bool = False, skip_wallet_balance_check: bool = False, ) -> Callable[..., tuple[bool, str]]: - created = datetime.now(tz=timezone.utc) + created = datetime.now(tz=UTC) def _condition( lm: ThlLedgerManager, diff --git a/generalresearch/managers/thl/ledger_manager/ledger.py b/generalresearch/managers/thl/ledger_manager/ledger.py index 864f1dd..00fac27 100644 --- a/generalresearch/managers/thl/ledger_manager/ledger.py +++ b/generalresearch/managers/thl/ledger_manager/ledger.py @@ -3,8 +3,9 @@ from __future__ import annotations import logging from collections import defaultdict from collections.abc import Collection -from datetime import datetime, timedelta, timezone -from typing import Any, Callable +from datetime import datetime, timedelta, timezone, UTC +from typing import Any +from collections.abc import Callable from uuid import UUID import redis @@ -104,8 +105,8 @@ class LedgerManagerBasePostgres(PostgresManager, RedisManager): filters = [] params = {} if time_start or time_end: - time_end = time_end or datetime.now(tz=timezone.utc) - time_start = time_start or datetime(2017, 1, 1, tzinfo=timezone.utc) + time_end = time_end or datetime.now(tz=UTC) + time_start = time_start or datetime(2017, 1, 1, tzinfo=UTC) assert time_start.tzinfo.utcoffset(time_start) == timedelta() assert time_end.tzinfo.utcoffset(time_end) == timedelta() filters.append("lt.created BETWEEN %(time_start)s AND %(time_end)s") @@ -152,7 +153,7 @@ class LedgerTransactionManager(LedgerManagerBasePostgres): if metadata is None: metadata = dict() if created is None: - created = datetime.now(tz=timezone.utc) + created = datetime.now(tz=UTC) t = LedgerTransaction( created=created, @@ -449,7 +450,7 @@ class LedgerTransactionManager(LedgerManagerBasePostgres): id=row["transaction_id"], entries=entries, metadata=metadata, - created=row["created"].replace(tzinfo=timezone.utc), + created=row["created"].replace(tzinfo=UTC), ext_description=row["ext_description"], tag=row["tag"], ) diff --git a/generalresearch/managers/thl/ledger_manager/thl_ledger.py b/generalresearch/managers/thl/ledger_manager/thl_ledger.py index 3119a46..c9203f6 100644 --- a/generalresearch/managers/thl/ledger_manager/thl_ledger.py +++ b/generalresearch/managers/thl/ledger_manager/thl_ledger.py @@ -2,9 +2,10 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from decimal import Decimal -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING +from collections.abc import Callable from uuid import UUID import numpy as np @@ -244,10 +245,10 @@ class ThlLedgerManager(LedgerManager): time_end: datetime | None = None, ): if time_start is None: - time_start = datetime(year=2017, month=1, day=1, tzinfo=timezone.utc) + time_start = datetime(year=2017, month=1, day=1, tzinfo=UTC) if time_end is None: - time_end = datetime.now(tz=timezone.utc) + time_end = datetime.now(tz=UTC) assert all( isinstance(item, str) for item in account_uuids @@ -798,7 +799,7 @@ class ThlLedgerManager(LedgerManager): skip_flag_check = True assert ( - datetime.now(tz=timezone.utc) > created + datetime.now(tz=UTC) > created ), "created cannot be in the future" f = lambda: self.create_tx_bp_payout_( product=product, @@ -904,7 +905,7 @@ class ThlLedgerManager(LedgerManager): for retry of a failed previous call. """ assert ( - datetime.now(tz=timezone.utc) > created + datetime.now(tz=UTC) > created ), "created cannot be in the future" assert isinstance(amount, int) assert isinstance(amount, USDCent) @@ -1837,7 +1838,7 @@ class ThlLedgerManager(LedgerManager): user.product.user_wallet_config.enabled ), "Can't get wallet balance on non-managed account." - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) wallet = self.get_account_or_create_user_wallet(user) if user.product_id == JAMES_BILLINGS_BPID: assert since_days_ago is None @@ -1867,7 +1868,7 @@ class ThlLedgerManager(LedgerManager): After 3 days, about 25% of all "future" recons have happened, 7 days: 50%, 14 days: 75%, till end of next month: 100%. """ - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) # The redeemable balance can NOT ever be more than the actual user_wallet_balance # Sum up the redeemable amount for each complete diff --git a/generalresearch/managers/thl/product.py b/generalresearch/managers/thl/product.py index 46280b8..3b92361 100644 --- a/generalresearch/managers/thl/product.py +++ b/generalresearch/managers/thl/product.py @@ -4,7 +4,7 @@ import json import logging import operator from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from decimal import Decimal from threading import Lock from typing import TYPE_CHECKING @@ -293,7 +293,7 @@ class ProductManager(PostgresManager): UserWalletConfig, ) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) # TODO: Add product_id, and possibly name uniqueness validation to the # pydantic model definition itself. The create manager doesn't need diff --git a/generalresearch/managers/thl/profiling/uqa.py b/generalresearch/managers/thl/profiling/uqa.py index cbe39e7..3854333 100644 --- a/generalresearch/managers/thl/profiling/uqa.py +++ b/generalresearch/managers/thl/profiling/uqa.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from generalresearch.managers.base import PostgresManagerWithRedis from generalresearch.models.thl.profiling.user_question_answer import ( @@ -128,7 +128,7 @@ class UQAManager(PostgresManagerWithRedis): def get_from_db(self, user: User) -> list[UserQuestionAnswer]: logger.info(f"get_uqa_from_db: {user.user_id}") # Only store the latest row per question_id. We don't need it multiple times. - since = datetime.now(tz=timezone.utc) - timedelta(days=30) + since = datetime.now(tz=UTC) - timedelta(days=30) # We CAN use the RR, b/c either # 1) the cache expired and the user hasn't sent an answer recently diff --git a/generalresearch/managers/thl/profiling/user_upk.py b/generalresearch/managers/thl/profiling/user_upk.py index 1820103..ee36124 100644 --- a/generalresearch/managers/thl/profiling/user_upk.py +++ b/generalresearch/managers/thl/profiling/user_upk.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from collections import defaultdict from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from typing import Any from uuid import UUID @@ -56,7 +56,7 @@ class UserUpkManager(PostgresManagerWithRedis): return res def get_user_upk_mysql(self, user_id: int) -> list[UpkQuestionAnswer]: - since = datetime.now(tz=timezone.utc) - timedelta(days=89) + since = datetime.now(tz=UTC) - timedelta(days=89) query = """ SELECT diff --git a/generalresearch/managers/thl/session.py b/generalresearch/managers/thl/session.py index 746a518..771f882 100644 --- a/generalresearch/managers/thl/session.py +++ b/generalresearch/managers/thl/session.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta, timezone, UTC from decimal import Decimal from typing import Any from uuid import UUID, uuid4 @@ -188,7 +188,7 @@ class SessionManager(PostgresManager): # validation errors. There doesn't seem to be a clean way of doing this. # model_copy with update doesn't trigger the validators, so we # re-run model_validate after - finished = finished if finished else datetime.now(tz=timezone.utc) + finished = finished if finished else datetime.now(tz=UTC) session.update( **{ "status": status, @@ -451,13 +451,13 @@ class SessionManager(PostgresManager): params = {} if started_before or started_after: - started_after = started_after or datetime(2017, 1, 1, tzinfo=timezone.utc) - started_before = started_before or datetime.now(tz=timezone.utc) + started_after = started_after or datetime(2017, 1, 1, tzinfo=UTC) + started_before = started_before or datetime.now(tz=UTC) assert ( - started_after.tzinfo == timezone.utc + started_after.tzinfo == UTC ), "started_after must be tz-aware as UTC" assert ( - started_before.tzinfo == timezone.utc + started_before.tzinfo == UTC ), "started_before must be tz-aware as UTC" assert ( started_after < started_before @@ -467,13 +467,13 @@ class SessionManager(PostgresManager): params["started_before"] = started_before if adjusted_before or adjusted_after: - adjusted_after = adjusted_after or datetime(2017, 1, 1, tzinfo=timezone.utc) - adjusted_before = adjusted_before or datetime.now(tz=timezone.utc) + adjusted_after = adjusted_after or datetime(2017, 1, 1, tzinfo=UTC) + adjusted_before = adjusted_before or datetime.now(tz=UTC) assert ( - adjusted_after.tzinfo == timezone.utc + adjusted_after.tzinfo == UTC ), "adjusted_after must be tz-aware as UTC" assert ( - adjusted_before.tzinfo == timezone.utc + adjusted_before.tzinfo == UTC ), "adjusted_before must be tz-aware as UTC" assert ( adjusted_after < adjusted_before diff --git a/generalresearch/managers/thl/survey.py b/generalresearch/managers/thl/survey.py index c7671ce..966e96d 100644 --- a/generalresearch/managers/thl/survey.py +++ b/generalresearch/managers/thl/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections import defaultdict from collections.abc import Collection -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import Any import pandas as pd @@ -544,7 +544,7 @@ class SurveyStatManager(PostgresManager): VALUES ({values_str}) ON CONFLICT ({unique_cols_str}) DO UPDATE SET {update_str};""" - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) params = [ss.model_dump_sql() | {"updated_at": now} for ss in survey_stats] with self.pg_config.make_connection() as conn: diff --git a/generalresearch/managers/thl/task_adjustment.py b/generalresearch/managers/thl/task_adjustment.py index e4736d4..3ec3d41 100644 --- a/generalresearch/managers/thl/task_adjustment.py +++ b/generalresearch/managers/thl/task_adjustment.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from functools import cached_property @@ -120,8 +120,8 @@ class TaskAdjustmentManager(PostgresManager): CHANGES/DELTAS as just communicated by the marketplace, not what the Wall's final adjusted_* will be. """ - alert_time = alert_time or datetime.now(tz=timezone.utc) - assert alert_time.tzinfo == timezone.utc + alert_time = alert_time or datetime.now(tz=UTC) + assert alert_time.tzinfo == UTC wall = self.wall_manager.get_from_uuid(wall_uuid) session = self.session_manager.get_from_id(wall.session_id) diff --git a/generalresearch/managers/thl/user_compensate.py b/generalresearch/managers/thl/user_compensate.py index 543de87..c6c0747 100644 --- a/generalresearch/managers/thl/user_compensate.py +++ b/generalresearch/managers/thl/user_compensate.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from uuid import uuid4 @@ -28,7 +28,7 @@ def user_compensate( pg_config = ledger_manager.pg_config redis_client = ledger_manager.redis_client - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) assert type(amount_int) is int user.prefetch_product(pg_config=pg_config) assert ( diff --git a/generalresearch/managers/thl/user_manager/mysql_user_manager.py b/generalresearch/managers/thl/user_manager/mysql_user_manager.py index dbed5de..7931ba4 100644 --- a/generalresearch/managers/thl/user_manager/mysql_user_manager.py +++ b/generalresearch/managers/thl/user_manager/mysql_user_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from functools import lru_cache from uuid import uuid4 @@ -26,7 +26,7 @@ class MysqlUserManager: def _set_last_seen(self, user: User) -> None: # Don't call this directly. Use UserManager.set_last_seen() assert not self.is_read_replica - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) self.pg_config.execute_write( """ UPDATE thl_user @@ -118,7 +118,7 @@ class MysqlUserManager: if not self.product_id_exists(product_id=product_id): raise ValueError(f"userprofile_brokerageproduct not found: {product_id}") - now = created or datetime.now(tz=timezone.utc) + now = created or datetime.now(tz=UTC) user_uuid = uuid4().hex params = { "user_uuid": user_uuid, diff --git a/generalresearch/managers/thl/userhealth.py b/generalresearch/managers/thl/userhealth.py index fe2163f..0bc60ec 100644 --- a/generalresearch/managers/thl/userhealth.py +++ b/generalresearch/managers/thl/userhealth.py @@ -2,7 +2,7 @@ from __future__ import annotations import ipaddress from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from itertools import zip_longest from typing import Any @@ -210,7 +210,7 @@ class IPRecordManager(PostgresManagerWithRedis): data = { "user_id": user_id, "ip": ipaddress.ip_address(ip).exploded, - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), } fips_cols = [ @@ -335,7 +335,7 @@ class AuditLogManager(PostgresManager): al = AuditLog.model_validate( { "user_id": user_id, - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), "level": level, "event_type": event_type, "event_msg": event_msg, @@ -495,7 +495,7 @@ class AuditLogManager(PostgresManager): ), "must pass user_id as int" if created_after is None: - created_after = datetime.now(tz=timezone.utc) - timedelta(days=7) + created_after = datetime.now(tz=UTC) - timedelta(days=7) filters = [ "user_id = ANY(%(user_ids)s)", diff --git a/generalresearch/managers/thl/wall.py b/generalresearch/managers/thl/wall.py index c2eb821..7e413d7 100644 --- a/generalresearch/managers/thl/wall.py +++ b/generalresearch/managers/thl/wall.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging from collections import defaultdict from collections.abc import Collection -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from functools import cached_property from uuid import uuid4 @@ -360,12 +360,10 @@ class WallManager(PostgresManager): params = {} filters.append("user_id = %(user_id)s") params["user_id"] = user_id - default_started = datetime.now(tz=timezone.utc) - timedelta(days=90) + default_started = datetime.now(tz=UTC) - timedelta(days=90) started_after = started_after or default_started - started_before = started_before or datetime.now(tz=timezone.utc) - assert ( - started_before.tzinfo == timezone.utc - ), "started_before must be tz-aware as UTC" + started_before = started_before or datetime.now(tz=UTC) + assert started_before.tzinfo == UTC, "started_before must be tz-aware as UTC" assert ( started_after < started_before ), "started_after must be before started_before" @@ -412,7 +410,7 @@ class WallManager(PostgresManager): started_before: datetime | None = None, order_by: str | None = "-started", ) -> list[WallAttempt]: - started_before = started_before or datetime.now(tz=timezone.utc) + started_before = started_before or datetime.now(tz=UTC) res = [] page = 1 while True: diff --git a/generalresearch/managers/thl/wallet/__init__.py b/generalresearch/managers/thl/wallet/__init__.py index b063e54..9f3ae85 100644 --- a/generalresearch/managers/thl/wallet/__init__.py +++ b/generalresearch/managers/thl/wallet/__init__.py @@ -32,8 +32,8 @@ def manage_pending_cashout( user_ip_history_manager: UserIpHistoryManager, user_manager: UserManager, ledger_manager: ThlLedgerManager, - order_data: Optional[Union[Dict[str, Any], CashMailOrderData]] = None, - tango_client: Optional[TangoClient] = None, + order_data: dict[str, Any] | CashMailOrderData | None = None, + tango_client: TangoClient | None = None, ) -> UserPayoutEvent: """ Called by a UI actions performed by Todd. This rejects/approves/cancels diff --git a/generalresearch/managers/thl/wallet/tango.py b/generalresearch/managers/thl/wallet/tango.py index 2f2dc52..445719d 100644 --- a/generalresearch/managers/thl/wallet/tango.py +++ b/generalresearch/managers/thl/wallet/tango.py @@ -65,8 +65,8 @@ def complete_tango_order( def create_tango_order( - request_data: Dict[str, Any], ref_id: str, tango_client: TangoClient -) -> Dict[str, Any]: + request_data: dict[str, Any], ref_id: str, tango_client: TangoClient +) -> dict[str, Any]: """ Create a tango gift card order. Throws exception if anything is not right. diff --git a/generalresearch/models/admin/__init__.py b/generalresearch/models/admin/__init__.py index ebe839a..ad6302b 100644 --- a/generalresearch/models/admin/__init__.py +++ b/generalresearch/models/admin/__init__.py @@ -1,14 +1,14 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone import pandas as pd from dateutil import relativedelta def get_date_list(start_datetime: datetime, end_datetime: datetime | None = None): - start_datetime = start_datetime.replace(tzinfo=timezone.utc) - end_datetime = end_datetime if end_datetime else datetime.now(tz=timezone.utc) + start_datetime = start_datetime.replace(tzinfo=UTC) + end_datetime = end_datetime if end_datetime else datetime.now(tz=UTC) return ( pd.date_range(start_datetime, end_datetime, freq="1D") .strftime("%Y-%m-%d") @@ -22,7 +22,7 @@ def year_start(periods_ago: int = 6) -> datetime: years. Goal is to provide a simple way to know when to do filters from """ - n: datetime = datetime.now(tz=timezone.utc) + n: datetime = datetime.now(tz=UTC) d: datetime = n - relativedelta.relativedelta(years=periods_ago) return d.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0) @@ -33,7 +33,7 @@ def month_start(periods_ago: int = 6) -> datetime: months. Goal is to provide a simple way to know when to do filters from """ - n: datetime = datetime.now(tz=timezone.utc) + n: datetime = datetime.now(tz=UTC) d: datetime = n - relativedelta.relativedelta(months=periods_ago) return d.replace(day=1, hour=0, minute=0, second=0, microsecond=0) @@ -44,7 +44,7 @@ def day_start(periods_ago: int = 6) -> datetime: days. Goal is to provide a simple way to know when to do filters from """ - n: datetime = datetime.now(tz=timezone.utc) + n: datetime = datetime.now(tz=UTC) d: datetime = n - relativedelta.relativedelta(days=periods_ago) return d.replace(hour=0, minute=0, second=0, microsecond=0) @@ -55,6 +55,6 @@ def hour_start(periods_ago: int = 6) -> datetime: hours. Goal is to provide a simple way to know when to do filters from """ - n: datetime = datetime.now(tz=timezone.utc) + n: datetime = datetime.now(tz=UTC) d: datetime = n - relativedelta.relativedelta(hours=periods_ago) return d.replace(minute=0, second=0, microsecond=0) diff --git a/generalresearch/models/admin/request.py b/generalresearch/models/admin/request.py index 67bd263..2d68de1 100644 --- a/generalresearch/models/admin/request.py +++ b/generalresearch/models/admin/request.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from enum import Enum from typing import Literal @@ -25,9 +25,9 @@ class ReportRequest(BaseModel): index1: str = Field(default="product_id") start: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - timedelta(days=14) + default_factory=lambda: datetime.now(tz=UTC) - timedelta(days=14) ) - end: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=timezone.utc)) + end: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) interval: Literal["5min", "15min", "1h", "6h", "12h", "1d"] = "1h" include_open_bucket: bool = Field(default=True) @@ -35,7 +35,7 @@ class ReportRequest(BaseModel): @computed_field( title="Start floor", description="The datetime that this report starts from", - examples=[datetime(year=2025, month=5, day=1, tzinfo=timezone.utc)], + examples=[datetime(year=2025, month=5, day=1, tzinfo=UTC)], return_type=datetime, ) @property @@ -60,7 +60,7 @@ class ReportRequest(BaseModel): @model_validator(mode="after") def check_start_end_tz(self): - assert self.start.tzinfo == self.end.tzinfo == timezone.utc + assert self.start.tzinfo == self.end.tzinfo == UTC return self @model_validator(mode="after") @@ -150,7 +150,7 @@ class ReportRequest(BaseModel): start=self.ts_start_floor, end=self.ts_end, freq=self.interval, - tz=timezone.utc, + tz=UTC, ) def bucket_ranges(self) -> list[tuple[pd.Timestamp, pd.Timestamp]]: diff --git a/generalresearch/models/cint/__init__.py b/generalresearch/models/cint/__init__.py index 2c1be7e..d2713ab 100644 --- a/generalresearch/models/cint/__init__.py +++ b/generalresearch/models/cint/__init__.py @@ -1,5 +1,6 @@ +from typing import Annotated + from pydantic import Field -from typing_extensions import Annotated CintQuestionIdType = Annotated[ str, Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") diff --git a/generalresearch/models/cint/question.py b/generalresearch/models/cint/question.py index 1ac9eea..5959141 100644 --- a/generalresearch/models/cint/question.py +++ b/generalresearch/models/cint/question.py @@ -1,13 +1,12 @@ from __future__ import annotations import json -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator -from typing_extensions import Self from generalresearch.models import Source, string_utils from generalresearch.models.cint import CintQuestionIdType @@ -151,7 +150,7 @@ class CintQuestion(MarketplaceQuestion): options = None created_at = datetime.strptime( d["create_date"], "%Y-%m-%dT%H:%M:%S%z" - ).astimezone(timezone.utc) + ).astimezone(UTC) if d.get("question_options"): options = [ @@ -189,7 +188,7 @@ class CintQuestion(MarketplaceQuestion): ] if d.get("created_at"): - d["created_at"] = d["created_at"].replace(tzinfo=timezone.utc) + d["created_at"] = d["created_at"].replace(tzinfo=UTC) return cls( question_id=d["question_id"], diff --git a/generalresearch/models/cint/survey.py b/generalresearch/models/cint/survey.py index 56384e3..01615f6 100644 --- a/generalresearch/models/cint/survey.py +++ b/generalresearch/models/cint/survey.py @@ -2,9 +2,9 @@ from __future__ import annotations import json import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal -from typing import Any, Literal, Type +from typing import Annotated, Any, Literal, Self, Type from more_itertools import flatten from pydantic import ( @@ -15,7 +15,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Annotated, Self from generalresearch.locales import Localelator from generalresearch.models import Source, TaskCalculationType @@ -291,7 +290,7 @@ class CintSurvey(MarketplaceTask): return data @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return CintCondition @property @@ -390,7 +389,7 @@ class CintSurvey(MarketplaceTask): d["conditions"][q.criterion_hash] = q d["quotas"] = quotas - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d["created_at"] = now d["last_updated"] = now @@ -418,8 +417,8 @@ class CintSurvey(MarketplaceTask): @classmethod def from_mysql(cls, d: Dict[str, Any]) -> Self: - d["created_at"] = d["created_at"].replace(tzinfo=timezone.utc) - d["last_updated"] = d["last_updated"].replace(tzinfo=timezone.utc) + d["created_at"] = d["created_at"].replace(tzinfo=UTC) + d["last_updated"] = d["last_updated"].replace(tzinfo=UTC) d["qualifications"] = json.loads(d["qualifications"]) d["used_question_ids"] = json.loads(d["used_question_ids"]) d["quotas"] = json.loads(d["quotas"]) diff --git a/generalresearch/models/cint/task_collection.py b/generalresearch/models/cint/task_collection.py index 5d39090..efb6e95 100644 --- a/generalresearch/models/cint/task_collection.py +++ b/generalresearch/models/cint/task_collection.py @@ -35,8 +35,8 @@ CintTaskCollectionSchema = DataFrameSchema( "bid_ir": Column(float, Check.between(0, 1), nullable=True), "created_at": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), "last_updated": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), - "used_question_ids": Column(List[str]), - "all_hashes": Column(List[str]), # set >> list for column support + "used_question_ids": Column(list[str]), + "all_hashes": Column(list[str]), # set >> list for column support }, checks=[], index=Index( diff --git a/generalresearch/models/custom_types.py b/generalresearch/models/custom_types.py index 84bf8e3..9346064 100644 --- a/generalresearch/models/custom_types.py +++ b/generalresearch/models/custom_types.py @@ -3,8 +3,8 @@ from __future__ import annotations import json import re import sys as _sys -from datetime import datetime, timedelta, timezone -from typing import Any, Literal +from datetime import UTC, datetime, timedelta, timezone +from typing import Annotated, Any, Literal from uuid import UUID from pydantic import ( @@ -20,7 +20,6 @@ from pydantic.functional_serializers import PlainSerializer from pydantic.functional_validators import AfterValidator, BeforeValidator from pydantic.networks import IPvAnyNetwork, UrlConstraints from pydantic_core import MultiHostHost, Url -from typing_extensions import Annotated from generalresearch.models import DeviceType, Source @@ -57,19 +56,17 @@ def convert_str_dt(v: Any) -> AwareDatetime | None: # to parse a str that was dumped using the iso8601 format with Z suffix. if v is not None and type(v) is str: assert v.endswith("Z") and "T" in v, "invalid format" - return datetime.strptime(v, "%Y-%m-%dT%H:%M:%S.%fZ").replace( - tzinfo=timezone.utc - ) + return datetime.strptime(v, "%Y-%m-%dT%H:%M:%S.%fZ").replace(tzinfo=UTC) return v def assert_utc(v: AwareDatetime) -> AwareDatetime: if isinstance(v, datetime): # We need utcoffset b/c FastAPI parses datetimes using FixedTimezone - assert v.tzinfo == timezone.utc or v.tzinfo.utcoffset(v) == timedelta( + assert v.tzinfo == UTC or v.tzinfo.utcoffset(v) == timedelta( 0 ), "Timezone is not UTC" - v = v.astimezone(timezone.utc) + v = v.astimezone(UTC) return v @@ -309,4 +306,4 @@ PropertyCode = Annotated[ def now_utc_factory(): - return datetime.now(tz=timezone.utc) + return datetime.now(tz=UTC) diff --git a/generalresearch/models/dynata/survey.py b/generalresearch/models/dynata/survey.py index 0e1b3e5..097eea2 100644 --- a/generalresearch/models/dynata/survey.py +++ b/generalresearch/models/dynata/survey.py @@ -2,10 +2,10 @@ from __future__ import annotations import json import logging -from datetime import timezone +from datetime import UTC, timezone from decimal import Decimal from functools import cached_property -from typing import Any, Literal, Type +from typing import Any, Literal, Self, Type from more_itertools import flatten from pydantic import ( @@ -17,7 +17,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import Source, TaskCalculationType @@ -500,7 +499,7 @@ class DynataSurvey(MarketplaceTask): return res @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return DynataCondition @property @@ -552,8 +551,8 @@ class DynataSurvey(MarketplaceTask): @classmethod def from_db(cls, d: Dict[str, Any]) -> Self: - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["last_updated"] = d["last_updated"].replace(tzinfo=timezone.utc) + d["created"] = d["created"].replace(tzinfo=UTC) + d["last_updated"] = d["last_updated"].replace(tzinfo=UTC) d["filters"] = json.loads(d["filters"]) d["quotas"] = json.loads(d["quotas"]) d["used_question_ids"] = json.loads(d["used_question_ids"]) diff --git a/generalresearch/models/events.py b/generalresearch/models/events.py index 63ed2a1..5efd4f6 100644 --- a/generalresearch/models/events.py +++ b/generalresearch/models/events.py @@ -1,6 +1,6 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from enum import StrEnum -from typing import Dict, Literal, Optional, Union +from typing import Annotated, Dict, Literal, Optional, Union from uuid import uuid4 from pydantic import ( @@ -12,7 +12,6 @@ from pydantic import ( TypeAdapter, model_validator, ) -from typing_extensions import Annotated from generalresearch.models import Source from generalresearch.models.custom_types import ( @@ -63,7 +62,7 @@ class TaskEnterPayload(BaseModel): source: Source = Field() survey_id: str = Field(min_length=1, max_length=32, examples=["127492892"]) - quota_id: Optional[str] = Field( + quota_id: str | None = Field( default=None, max_length=32, description="The marketplace's internal quota id", @@ -76,9 +75,9 @@ class TaskFinishPayload(TaskEnterPayload): duration_sec: PositiveFloat = Field() status: Status - status_code_1: Optional[StatusCode1] = None - status_code_2: Optional[WallStatusCode2] = None - cpi: Optional[NonNegativeInt] = Field(le=4000, default=None) + status_code_1: StatusCode1 | None = None + status_code_2: WallStatusCode2 | None = None + cpi: NonNegativeInt | None = Field(le=4000, default=None) class SessionEnterPayload(BaseModel): @@ -91,18 +90,13 @@ class SessionFinishPayload(SessionEnterPayload): duration_sec: PositiveFloat = Field() status: Status - status_code_1: Optional[StatusCode1] = None - status_code_2: Optional[SessionStatusCode2] = None - user_payout: Optional[NonNegativeInt] = Field(default=None, le=4000, ge=0) + status_code_1: StatusCode1 | None = None + status_code_2: SessionStatusCode2 | None = None + user_payout: NonNegativeInt | None = Field(default=None, le=4000, ge=0) EventPayload = Annotated[ - Union[ - TaskEnterPayload, - TaskFinishPayload, - SessionEnterPayload, - SessionFinishPayload, - ], + TaskEnterPayload | TaskFinishPayload | SessionEnterPayload | SessionFinishPayload, Field(discriminator="event_type"), ] @@ -110,12 +104,10 @@ EventPayload = Annotated[ class EventEnvelope(BaseModel): event_uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) event_type: EventType = Field() - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) version: int = 1 - product_user_id: Optional[str] = Field( + product_user_id: str | None = Field( min_length=3, max_length=128, examples=["app-user-9329ebd"], @@ -136,7 +128,7 @@ class EventEnvelope(BaseModel): class AggregateBySource(BaseModel): total: NonNegativeInt = Field(default=0) - by_source: Dict[Source, NonNegativeInt] = Field(default_factory=dict) + by_source: dict[Source, NonNegativeInt] = Field(default_factory=dict) @model_validator(mode="after") def remove_zero(self): @@ -145,8 +137,8 @@ class AggregateBySource(BaseModel): class MaxGaugeBySource(BaseModel): - value: Optional[NonNegativeInt] = Field(default=None) - by_source: Dict[Source, NonNegativeInt] = Field(default_factory=dict) + value: NonNegativeInt | None = Field(default=None) + by_source: dict[Source, NonNegativeInt] = Field(default_factory=dict) @model_validator(mode="after") def remove_zero(self): @@ -174,11 +166,9 @@ class StatsSnapshot(TaskStatsSnapshot): model_config = ConfigDict(ser_json_timedelta="float") # If this is set, then everything is scoped to this country. - country_iso: Optional[CountryISOLike] = Field(default=None) + country_iso: CountryISOLike | None = Field(default=None) - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # Counts: User related active_users_last_1h: NonNegativeInt = Field( @@ -217,17 +207,17 @@ class StatsSnapshot(TaskStatsSnapshot): ) # Rolling averages - session_avg_payout_last_24h: Optional[NonNegativeInt] = Field( + session_avg_payout_last_24h: NonNegativeInt | None = Field( description="Average (actual) payout of all tasks completed in the past 24 hrs" ) - session_avg_user_payout_last_24h: Optional[NonNegativeInt] = Field( + session_avg_user_payout_last_24h: NonNegativeInt | None = Field( description="Average (actual) user payout of all tasks completed in the past 24 hrs" ) - session_fail_avg_loi_last_24h: Optional[timedelta] = Field( + session_fail_avg_loi_last_24h: timedelta | None = Field( description="Average LOI of all tasks terminated in the past 24 hrs (excludes abandons)" ) - session_complete_avg_loi_last_24h: Optional[timedelta] = Field( + session_complete_avg_loi_last_24h: timedelta | None = Field( description="Average LOI of all tasks completed in the past 24 hrs" ) @@ -246,34 +236,26 @@ class StatsSnapshot(TaskStatsSnapshot): class EventMessage(BaseModel): kind: Literal[MessageKind.EVENT] = Field(default=MessageKind.EVENT) - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) data: EventEnvelope class StatsMessage(BaseModel): kind: Literal[MessageKind.STATS] = Field(default=MessageKind.STATS) - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # The data/StatsSnapshot can optionally be scoped to a country - country_iso: Optional[CountryISOLike] = Field(default=None) + country_iso: CountryISOLike | None = Field(default=None) data: StatsSnapshot class PingMessage(BaseModel): kind: Literal[MessageKind.PING] = Field(default=MessageKind.PING) - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) class PongMessage(BaseModel): kind: Literal[MessageKind.PONG] = Field(default=MessageKind.PONG) - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) class SubscribeMessage(BaseModel): diff --git a/generalresearch/models/gr/authentication.py b/generalresearch/models/gr/authentication.py index 4ee70f9..67a8fc2 100644 --- a/generalresearch/models/gr/authentication.py +++ b/generalresearch/models/gr/authentication.py @@ -3,8 +3,8 @@ from __future__ import annotations import binascii import json import os -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any +from datetime import UTC, datetime, timezone +from typing import TYPE_CHECKING, Any, Self from pydantic import ( AnyHttpUrl, @@ -15,7 +15,6 @@ from pydantic import ( PositiveInt, field_validator, ) -from typing_extensions import Self from generalresearch.decorators import LOG from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr @@ -199,7 +198,7 @@ class GRUser(BaseModel): @field_validator("date_joined") @classmethod def date_joined_utc(cls, v: datetime) -> datetime: - return v.replace(tzinfo=timezone.utc) + return v.replace(tzinfo=UTC) # --- Properties --- @property @@ -290,7 +289,7 @@ class GRUser(BaseModel): @classmethod def from_postgresql(cls, d: dict) -> Self: - d["date_joined"] = d["date_joined"].replace(tzinfo=timezone.utc) + d["date_joined"] = d["date_joined"].replace(tzinfo=UTC) return GRUser.model_validate(d) @classmethod @@ -354,7 +353,7 @@ class GRToken(BaseModel): @field_validator("created", mode="before") @classmethod def created_utc(cls, v: datetime) -> datetime: - return v.replace(tzinfo=timezone.utc) + return v.replace(tzinfo=UTC) # --- Properties --- diff --git a/generalresearch/models/gr/business.py b/generalresearch/models/gr/business.py index 51317da..a67cf48 100644 --- a/generalresearch/models/gr/business.py +++ b/generalresearch/models/gr/business.py @@ -3,10 +3,10 @@ from __future__ import annotations import json import logging import os -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from uuid import uuid4 import pandas as pd @@ -16,7 +16,6 @@ from psycopg.rows import dict_row from pydantic import BaseModel, ConfigDict, Field, PositiveInt from pydantic.json_schema import SkipJsonSchema from pydantic_extra_types.phone_numbers import PhoneNumber -from typing_extensions import Self from generalresearch.currency import USDCent from generalresearch.decorators import LOG @@ -391,8 +390,8 @@ class Business(BaseModel): pop_ledger = plm(ds=ds) if at_timestamp is None: - at_timestamp = datetime.now(tz=timezone.utc) - assert at_timestamp.tzinfo == timezone.utc + at_timestamp = datetime.now(tz=UTC) + assert at_timestamp.tzinfo == UTC ddf = pop_ledger.ddf( force_rr_latest=False, @@ -724,7 +723,7 @@ class Business(BaseModel): if "pop_financial" in keys: # We should explicitly pass the pop_financial years we want. By default, # at least get this year. - year = datetime.now(tz=timezone.utc).year + year = datetime.now(tz=UTC).year keys = list(set(keys) | {f"pop_financial:{year}"}) rc = gr_redis_config.create_redis_client() diff --git a/generalresearch/models/gr/team.py b/generalresearch/models/gr/team.py index 8d60825..900062f 100644 --- a/generalresearch/models/gr/team.py +++ b/generalresearch/models/gr/team.py @@ -2,10 +2,10 @@ from __future__ import annotations import json import os -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Self from uuid import uuid4 import pandas as pd @@ -18,7 +18,6 @@ from pydantic import ( field_validator, ) from pydantic.json_schema import SkipJsonSchema -from typing_extensions import Self from generalresearch.decorators import LOG from generalresearch.incite.mergers.foundations.enriched_session import ( @@ -92,7 +91,7 @@ class Membership(BaseModel): @classmethod def created_utc(cls, v: datetime | str) -> datetime | str: if isinstance(v, datetime): - return v.replace(tzinfo=timezone.utc) + return v.replace(tzinfo=UTC) return v # --- prefetch methods --- diff --git a/generalresearch/models/innovate/__init__.py b/generalresearch/models/innovate/__init__.py index 054c69d..26946e9 100644 --- a/generalresearch/models/innovate/__init__.py +++ b/generalresearch/models/innovate/__init__.py @@ -1,7 +1,7 @@ from enum import Enum +from typing import Annotated from pydantic import StringConstraints -from typing_extensions import Annotated # Note, this is called the KEY in the Question model InnovateQuestionID = Annotated[ diff --git a/generalresearch/models/innovate/survey.py b/generalresearch/models/innovate/survey.py index bcd50d3..0359bd6 100644 --- a/generalresearch/models/innovate/survey.py +++ b/generalresearch/models/innovate/survey.py @@ -2,13 +2,14 @@ from __future__ import annotations import json import logging -from datetime import date, timezone +from datetime import UTC, date, timezone from decimal import Decimal from functools import cached_property from typing import ( Annotated, Any, Literal, + Self, Type, ) @@ -20,7 +21,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import ( @@ -290,7 +290,7 @@ class InnovateSurvey(MarketplaceTask): return cls.model_validate(d) @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return InnovateCondition @property @@ -361,10 +361,10 @@ class InnovateSurvey(MarketplaceTask): @classmethod def from_db(cls, d: dict[str, Any]) -> Self: - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) - d["modified_api"] = d["modified_api"].replace(tzinfo=timezone.utc) - d["created_api"] = d["created_api"].replace(tzinfo=timezone.utc) + d["created"] = d["created"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) + d["modified_api"] = d["modified_api"].replace(tzinfo=UTC) + d["created_api"] = d["created_api"].replace(tzinfo=UTC) d["qualifications"] = json.loads(d["qualifications"]) d["used_question_ids"] = json.loads(d["used_question_ids"]) d["quotas"] = json.loads(d["quotas"]) diff --git a/generalresearch/models/legacy/bucket.py b/generalresearch/models/legacy/bucket.py index 2650b0b..812241d 100644 --- a/generalresearch/models/legacy/bucket.py +++ b/generalresearch/models/legacy/bucket.py @@ -4,7 +4,7 @@ import logging import math from datetime import timedelta from decimal import Decimal -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import ( BaseModel, @@ -14,7 +14,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.models import Source from generalresearch.models.custom_types import ( diff --git a/generalresearch/models/legacy/questions.py b/generalresearch/models/legacy/questions.py index 81e794c..8e19e57 100644 --- a/generalresearch/models/legacy/questions.py +++ b/generalresearch/models/legacy/questions.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Annotated, Any, Self from pydantic import ( BaseModel, @@ -14,7 +14,6 @@ from pydantic import ( model_validator, ) from sentry_sdk import capture_exception -from typing_extensions import Annotated, Self from generalresearch.models.custom_types import UUIDStr from generalresearch.models.legacy.api_status import StatusResponse diff --git a/generalresearch/models/lucid/__init__.py b/generalresearch/models/lucid/__init__.py index c3365db..4653339 100644 --- a/generalresearch/models/lucid/__init__.py +++ b/generalresearch/models/lucid/__init__.py @@ -1,5 +1,6 @@ +from typing import Annotated + from pydantic import Field -from typing_extensions import Annotated LucidQuestionIdType = Annotated[ str, Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") diff --git a/generalresearch/models/lucid/question.py b/generalresearch/models/lucid/question.py index 908ce70..6cbcb73 100644 --- a/generalresearch/models/lucid/question.py +++ b/generalresearch/models/lucid/question.py @@ -2,10 +2,9 @@ from __future__ import annotations import logging from enum import Enum -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self from pydantic import BaseModel, Field, field_validator, model_validator -from typing_extensions import Self from generalresearch.models import Source from generalresearch.models.lucid import LucidQuestionIdType diff --git a/generalresearch/models/marketplace/summary.py b/generalresearch/models/marketplace/summary.py index f75c530..9551417 100644 --- a/generalresearch/models/marketplace/summary.py +++ b/generalresearch/models/marketplace/summary.py @@ -2,11 +2,10 @@ from __future__ import annotations from abc import ABC from collections.abc import Collection -from typing import Literal +from typing import Literal, Self import numpy as np from pydantic import BaseModel, ConfigDict, Field, computed_field -from typing_extensions import Self from generalresearch.models.thl.stats import StatisticalSummary diff --git a/generalresearch/models/morning/__init__.py b/generalresearch/models/morning/__init__.py index 2c61c49..1bc15a7 100644 --- a/generalresearch/models/morning/__init__.py +++ b/generalresearch/models/morning/__init__.py @@ -1,7 +1,7 @@ from enum import Enum +from typing import Annotated from pydantic import StringConstraints -from typing_extensions import Annotated # This is text-based, in lowercase. e.g. 'age', 'household_income' MorningQuestionID = Annotated[ diff --git a/generalresearch/models/morning/question.py b/generalresearch/models/morning/question.py index 0ab5030..8a1f729 100644 --- a/generalresearch/models/morning/question.py +++ b/generalresearch/models/morning/question.py @@ -1,10 +1,9 @@ import json from enum import Enum -from typing import Any, Literal, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional, Self from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import Source @@ -54,7 +53,7 @@ class MorningQuestionType(str, Enum): class MorningUserQuestionAnswer(MarketplaceUserQuestionAnswer): question_id: MorningQuestionID = Field() - question_type: Optional[MorningQuestionType] = Field(default=None) + question_type: MorningQuestionType | None = Field(default=None) # Did this answer come from us asking, or was it passed back from the # marketplace? Note, morning doesn't "pass back" answers, but we can # retrieve a user's profile through API, so it is possible to populate @@ -92,7 +91,7 @@ class MorningQuestion(MarketplaceQuestion): frozen=True, ) # API calls this "responses", but I think that is a confusing name - options: Optional[List[MorningQuestionOption]] = Field( + options: list[MorningQuestionOption] | None = Field( default=None, min_length=1, frozen=True ) @@ -119,7 +118,7 @@ class MorningQuestion(MarketplaceQuestion): return options @classmethod - def from_api(cls, d: Dict[str, Any], country_iso: str, language_iso: str): + def from_api(cls, d: dict[str, Any], country_iso: str, language_iso: str): options = None if d.get("responses"): options = [ @@ -138,7 +137,7 @@ class MorningQuestion(MarketplaceQuestion): ) @classmethod - def from_db(cls, d: Dict[str, Any]) -> Self: + def from_db(cls, d: dict[str, Any]) -> Self: options = None if d["options"]: options = [ @@ -162,7 +161,7 @@ class MorningQuestion(MarketplaceQuestion): ), ) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", by_alias=True) d["options"] = json.dumps(d["options"]) return d diff --git a/generalresearch/models/morning/survey.py b/generalresearch/models/morning/survey.py index 255a63c..3c5a0e0 100644 --- a/generalresearch/models/morning/survey.py +++ b/generalresearch/models/morning/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import logging -from datetime import timezone +from datetime import UTC, timezone from decimal import Decimal from functools import cached_property from typing import ( @@ -12,6 +12,7 @@ from typing import ( List, Literal, Optional, + Self, Set, Tuple, Type, @@ -27,7 +28,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import Source @@ -79,14 +79,14 @@ class MorningStatistics(BaseModel): # bid bid_loi: int = Field(validation_alias="estimated_length_of_interview", # le=120 * 60) # If num_completes == 0 , this gets returned as 0. it should be None - obs_median_loi: Optional[NonNegativeInt] = Field( + obs_median_loi: NonNegativeInt | None = Field( validation_alias="median_length_of_interview", default=None, le=120 * 60 ) # API returns 100 until 5 completes! Should be None. # This is calculated as the total completes divided by the total number of # finished sessions that passed the prescreener. - qualified_conversion: Optional[float] = Field( + qualified_conversion: float | None = Field( ge=0, le=1, description="conversion rate of qualified respondents" ) @@ -142,7 +142,7 @@ class MorningTaskStatistics(MorningStatistics): # relevant to quotas. # API returns 100 until 5 completes! Should be None ... - system_conversion: Optional[float] = Field( + system_conversion: float | None = Field( description="conversion rate of the system. completes divided by total number of entrants to the system", ge=0, le=1, @@ -166,8 +166,8 @@ class MorningTaskStatistics(MorningStatistics): class MorningCondition(MarketplaceCondition): model_config = ConfigDict(populate_by_name=True, frozen=False, extra="ignore") - question_id: Optional[MorningQuestionID] = Field(validation_alias="id") - values: List[Annotated[str, Field(max_length=128)]] = Field( + question_id: MorningQuestionID | None = Field(validation_alias="id") + values: list[Annotated[str, Field(max_length=128)]] = Field( validation_alias="response_ids" ) value_type: ConditionValueType = Field(default=ConditionValueType.LIST) @@ -184,11 +184,11 @@ class MorningQuota(MorningStatistics, MarketplaceTask): max_digits=5, validation_alias="cost_per_interview", ) - condition_hashes: List[str] = Field(min_length=1, default_factory=list) + condition_hashes: list[str] = Field(min_length=1, default_factory=list) # since the Quota is the MarketplaceTask, it needs these fields, copied from the Bid source: Literal[Source.MORNING_CONSULT] = Field(default=Source.MORNING_CONSULT) - used_question_ids: Set[MorningQuestionID] = Field(default_factory=set) + used_question_ids: set[MorningQuestionID] = Field(default_factory=set) country_iso: CountryISO = Field(frozen=True) country_isos: CountryISOs = Field() language_isos: LanguageISOs = Field(frozen=True) @@ -219,11 +219,11 @@ class MorningQuota(MorningStatistics, MarketplaceTask): @computed_field @cached_property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: return set(self.condition_hashes) @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return MorningCondition @property @@ -233,7 +233,7 @@ class MorningQuota(MorningStatistics, MarketplaceTask): @property def marketplace_genders( self, - ) -> Dict[Gender, Optional[MarketplaceCondition]]: + ) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: MorningCondition( question_id="gender", @@ -253,14 +253,14 @@ class MorningQuota(MorningStatistics, MarketplaceTask): # num_available includes in-progress (they're already deducted) return self.num_available >= self._min_open_spots - def passes(self, criteria_evaluation: Dict[str, Optional[bool]]) -> bool: + def passes(self, criteria_evaluation: dict[str, bool | None]) -> bool: # Passes means we 1) meet all conditions (aka "match") AND 2) the quota is open. return self.is_open and self.matches(criteria_evaluation) # TODO: I did some speed tests. This is faster than how this is implemented # in sago/spectrum/dynata/etc. We should generalize this logic instead of # copying/pasting it 7 times. (matches, matches_optional and _soft) - def matches(self, criteria_evaluation: Dict[str, Optional[bool]]) -> bool: + def matches(self, criteria_evaluation: dict[str, bool | None]) -> bool: # Matches means we meet all conditions. # In Morning, all quotas are mutually exclusive. so if it doesn't # matter if we match a closed quota, b/c that means that we won't @@ -268,8 +268,8 @@ class MorningQuota(MorningStatistics, MarketplaceTask): return self.matches_optional(criteria_evaluation) is True def matches_optional( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Optional[bool]: + self, criteria_evaluation: dict[str, bool | None] + ) -> bool | None: for c in self.condition_hashes: eval_value = criteria_evaluation.get(c) if eval_value is False: @@ -279,8 +279,8 @@ class MorningQuota(MorningStatistics, MarketplaceTask): return True def matches_soft( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], List[str]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, list[str]]: # Passes back "matches" (T/F/none) and a list of unknown criterion hashes unknowns = list() for c in self.condition_hashes: @@ -321,22 +321,22 @@ class MorningBid(MorningTaskStatistics): timeout: PositiveInt = Field(le=24 * 60 * 60) topic_id: str = Field(min_length=1, max_length=64) - exclusions: List[MorningExclusion] = Field(default_factory=list) + exclusions: list[MorningExclusion] = Field(default_factory=list) - quotas: List[MorningQuota] = Field(default_factory=list) + quotas: list[MorningQuota] = Field(default_factory=list) source: Literal[Source.MORNING_CONSULT] = Field(default=Source.MORNING_CONSULT) - used_question_ids: Set[MorningQuestionID] = Field(default_factory=set) + used_question_ids: set[MorningQuestionID] = Field(default_factory=set) # This is a "special" key to store all conditions that are used (as # "condition_hashes") throughout this survey. In the reduced representation # of this task (nearly always, for db i/o, in global_vars) this field will # be null. - conditions: Optional[Dict[str, MorningCondition]] = Field(default=None) + conditions: dict[str, MorningCondition] | None = Field(default=None) # This doesn't get stored in the db directly - experimental_single_use_qualifications: Optional[List[MorningQuestion]] = Field( + experimental_single_use_qualifications: list[MorningQuestion] | None = Field( default=None ) @@ -345,8 +345,8 @@ class MorningBid(MorningTaskStatistics): created_api: AwareDatetimeISO = Field(validation_alias="published_at") # This does not come from the API. We set it when we update this in the db. - created: Optional[AwareDatetimeISO] = Field(default=None) - updated: Optional[AwareDatetimeISO] = Field(default=None) + created: AwareDatetimeISO | None = Field(default=None) + updated: AwareDatetimeISO | None = Field(default=None) # ignoring from API: closed_at @@ -373,7 +373,7 @@ class MorningBid(MorningTaskStatistics): @computed_field @cached_property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set() for q in self.quotas: s.update(set(q.condition_hashes)) @@ -387,7 +387,7 @@ class MorningBid(MorningTaskStatistics): @model_validator(mode="before") @classmethod - def setup_quota_fields(cls, data: Dict[str, Any]) -> Dict[str, Any]: + def setup_quota_fields(cls, data: dict[str, Any]) -> dict[str, Any]: # These fields get "inherited" by each quota from its bid. quota_fields = [ "country_iso", @@ -419,7 +419,7 @@ class MorningBid(MorningTaskStatistics): @model_validator(mode="before") @classmethod - def setup_conditions(cls, data: Dict[str, Any]) -> Dict[str, Any]: + def setup_conditions(cls, data: dict[str, Any]) -> dict[str, Any]: if "conditions" in data: return data @@ -448,7 +448,7 @@ class MorningBid(MorningTaskStatistics): @model_validator(mode="before") @classmethod - def clean_alias(cls, data: Dict[str, Any]) -> Dict[str, Any]: + def clean_alias(cls, data: dict[str, Any]) -> dict[str, Any]: # Make sure fields are named certain ways, so we don't have to check # aliases within other validators if "estimated_length_of_interview" in data: @@ -503,18 +503,16 @@ class MorningBid(MorningTaskStatistics): return d @classmethod - def from_db(cls, d: Dict[str, Any]) -> Self: - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) - d["expected_end"] = d["expected_end"].replace(tzinfo=timezone.utc) - d["created_api"] = d["created_api"].replace(tzinfo=timezone.utc) + def from_db(cls, d: dict[str, Any]) -> Self: + d["created"] = d["created"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) + d["expected_end"] = d["expected_end"].replace(tzinfo=UTC) + d["created_api"] = d["created_api"].replace(tzinfo=UTC) d["used_question_ids"] = json.loads(d["used_question_ids"]) d["exclusions"] = json.loads(d["exclusions"]) return cls.model_validate(d) - def passes_quotas( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Optional[str]: + def passes_quotas(self, criteria_evaluation: dict[str, bool | None]) -> str | None: # Quotas are mutually-exclusive. A user can only possibly match 1 quota. # Returns the passing quota ID or None (if user doesn't pass any quota) for q in self.quotas: @@ -522,8 +520,8 @@ class MorningBid(MorningTaskStatistics): return q.id def passes_quotas_soft( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], Optional[List[str]], Optional[Set[str]]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, list[str] | None, set[str] | None]: """ Quotas are mutually-exclusive. A user can only possibly match 1 quota. As such, all unknown questions on any quota will be @@ -547,15 +545,15 @@ class MorningBid(MorningTaskStatistics): return False, None, None def determine_eligibility( - self, criteria_evaluation: dict[str, Optional[bool]] - ) -> Optional[str]: + self, criteria_evaluation: dict[str, bool | None] + ) -> str | None: if not self.is_open: return None return self.passes_quotas(criteria_evaluation) def determine_eligibility_soft( - self, criteria_evaluation: dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], Optional[List[str]], Optional[Set[str]]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, list[str] | None, set[str] | None]: if not self.is_open: return False, None, None return self.passes_quotas_soft(criteria_evaluation) diff --git a/generalresearch/models/network/mtr/execute.py b/generalresearch/models/network/mtr/execute.py index d77e814..953124d 100644 --- a/generalresearch/models/network/mtr/execute.py +++ b/generalresearch/models/network/mtr/execute.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from uuid import uuid4 from generalresearch.models.custom_types import UUIDStr @@ -33,10 +33,10 @@ def execute_mtr( ), ) - started_at = datetime.now(tz=timezone.utc) + started_at = datetime.now(tz=UTC) tool_version = get_mtr_version() result = run_mtr(config) - finished_at = datetime.now(tz=timezone.utc) + finished_at = datetime.now(tz=UTC) return MTRRun( tool_name=ToolName.MTR, diff --git a/generalresearch/models/network/nmap/parser.py b/generalresearch/models/network/nmap/parser.py index e946e5f..6ad4ab4 100644 --- a/generalresearch/models/network/nmap/parser.py +++ b/generalresearch/models/network/nmap/parser.py @@ -1,7 +1,7 @@ from __future__ import annotations import xml.etree.ElementTree as ET -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import Any from generalresearch.models.network.definitions import IPProtocol @@ -123,7 +123,7 @@ class NmapXmlParser: finished_at = None ts = finished.attrib.get("time") if ts: - finished_at = datetime.fromtimestamp(int(ts), tz=timezone.utc) + finished_at = datetime.fromtimestamp(int(ts), tz=UTC) return { "finished_at": finished_at, @@ -136,7 +136,7 @@ class NmapXmlParser: nmaprun = dict(nmaprun_el.attrib) nmap_data["command_line"] = nmaprun["args"] nmap_data["started_at"] = datetime.fromtimestamp( - float(nmaprun["start"]), tz=timezone.utc + float(nmaprun["start"]), tz=UTC ) nmap_data["version"] = nmaprun["version"] nmap_data["xmloutputversion"] = nmaprun["xmloutputversion"] diff --git a/generalresearch/models/network/nmap/result.py b/generalresearch/models/network/nmap/result.py index 3f9cae6..e6a0fd3 100644 --- a/generalresearch/models/network/nmap/result.py +++ b/generalresearch/models/network/nmap/result.py @@ -256,13 +256,13 @@ class NmapScanInfo(BaseModel): services: str = Field() @cached_property - def port_set(self) -> Set[int]: + def port_set(self) -> set[int]: """ Expand the Nmap services string into a set of port numbers. Example: "22-25,80,443" -> {22,23,24,25,80,443} """ - ports: Set[int] = set() + ports: set[int] = set() for part in self.services.split(","): if "-" in part: start, end = part.split("-", 1) diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py index cabd13c..1d74df2 100644 --- a/generalresearch/models/network/rdns/execute.py +++ b/generalresearch/models/network/rdns/execute.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from uuid import uuid4 from generalresearch.models.custom_types import UUIDStr @@ -21,11 +21,11 @@ from generalresearch.models.network.tool_run_command import ( def execute_rdns(ip: str, scan_group_id: UUIDStr | None = None): - started_at = datetime.now(tz=timezone.utc) + started_at = datetime.now(tz=UTC) tool_version = get_dig_version() config = RDNSRunCommand(options=RDNSRunCommandOptions(ip=ip)) result = run_rdns(config) - finished_at = datetime.now(tz=timezone.utc) + finished_at = datetime.now(tz=UTC) run = RDNSRun( tool_name=ToolName.DIG, diff --git a/generalresearch/models/precision/__init__.py b/generalresearch/models/precision/__init__.py index 4bb2c6a..089c34a 100644 --- a/generalresearch/models/precision/__init__.py +++ b/generalresearch/models/precision/__init__.py @@ -1,7 +1,7 @@ from enum import Enum +from typing import Annotated from pydantic import StringConstraints -from typing_extensions import Annotated class PrecisionStatus(str, Enum): diff --git a/generalresearch/models/precision/survey.py b/generalresearch/models/precision/survey.py index 646d60e..be98a79 100644 --- a/generalresearch/models/precision/survey.py +++ b/generalresearch/models/precision/survey.py @@ -1,9 +1,9 @@ from __future__ import annotations import json -from datetime import timezone +from datetime import UTC, timezone from functools import cached_property -from typing import Any, Dict, List, Literal, Optional, Self, Set, Tuple, Type +from typing import Annotated, Any, Dict, List, Literal, Optional, Self, Set, Tuple, Type from more_itertools import flatten from pydantic import ( @@ -14,7 +14,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Annotated from generalresearch.models import Source from generalresearch.models.custom_types import ( @@ -34,8 +33,8 @@ from generalresearch.models.thl.survey.condition import ( class PrecisionCondition(MarketplaceCondition): - question_id: Optional[PrecisionQuestionID] = Field() - values: List[Annotated[str, Field(max_length=128)]] = Field() + question_id: PrecisionQuestionID | None = Field() + values: list[Annotated[str, Field(max_length=128)]] = Field() value_type: ConditionValueType = Field(default=ConditionValueType.LIST) _CONVERT_LIST_TO_RANGE = ["age"] @@ -54,7 +53,7 @@ class PrecisionQuota(BaseModel): termination_count: int = Field(ge=0) overquota_count: int = Field(ge=0) - condition_hashes: List[str] = Field(min_length=1, default_factory=list) + condition_hashes: list[str] = Field(min_length=1, default_factory=list) # Min spots a quota should have open to be OPEN _min_open_spots: int = PrivateAttr(default=3) @@ -78,7 +77,7 @@ class PrecisionQuota(BaseModel): # TODO: I did some speed tests. This is faster than how this is implemented # in sago/spectrum/dynata/etc. We should generalize this logic instead of # copying/pasting it 7 times. (matches, matches_optional and _soft) - def matches(self, criteria_evaluation: Dict[str, Optional[bool]]) -> bool: + def matches(self, criteria_evaluation: dict[str, bool | None]) -> bool: # Matches means we meet all conditions. # In Morning, all quotas are mutually exclusive. so if it doesn't # matter if we match a closed quota, b/c that means that we won't @@ -86,8 +85,8 @@ class PrecisionQuota(BaseModel): return self.matches_optional(criteria_evaluation) is True def matches_optional( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Optional[bool]: + self, criteria_evaluation: dict[str, bool | None] + ) -> bool | None: for c in self.condition_hashes: eval_value = criteria_evaluation.get(c) if eval_value is False: @@ -97,8 +96,8 @@ class PrecisionQuota(BaseModel): return True def matches_soft( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], List[str]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, list[str]]: # Passes back "matches" (T/F/none) and a list of unknown criterion hashes unknowns = list() for c in self.condition_hashes: @@ -128,7 +127,7 @@ class PrecisionSurvey(MarketplaceTask): name: str = Field(validation_alias="prj_name") survey_guid: UUIDStrCoerce = Field(validation_alias="prj_guid") - category_id: Optional[str] = Field(validation_alias="sc_id", default=None) + category_id: str | None = Field(validation_alias="sc_id", default=None) buyer_id: CoercedStr = Field(max_length=16) # This seems to always be 0 ... ? @@ -141,7 +140,7 @@ class PrecisionSurvey(MarketplaceTask): bid_ir: float = Field(ge=0, le=1, validation_alias="ir") # Be careful with this, it doesn't make any sense. See survey 452481, has 12 completes with a 100% live_ir, # but the only quotas have 0 completes and 1052 terms. .... ?? - global_conversion: Optional[float] = Field( + global_conversion: float | None = Field( ge=0, le=1, default=None, @@ -156,31 +155,31 @@ class PrecisionSurvey(MarketplaceTask): allowed_devices: DeviceTypes = Field(min_length=1) entry_link: str = Field(validation_alias="url") - excluded_surveys: Optional[AlphaNumStrSet] = Field( + excluded_surveys: AlphaNumStrSet | None = Field( description="list of excluded survey ids", default=None, validation_alias="exclusion_project_id", ) - quotas: List[PrecisionQuota] = Field(default_factory=list) + quotas: list[PrecisionQuota] = Field(default_factory=list) source: Literal[Source.PRECISION] = Field(default=Source.PRECISION) - used_question_ids: Set[PrecisionQuestionID] = Field(default_factory=set) + used_question_ids: set[PrecisionQuestionID] = Field(default_factory=set) # This is a "special" key to store all conditions that are used (as "condition_hashes") throughout # this survey. In the reduced representation of this task (nearly always, for db i/o, in global_vars) # this field will be null. - conditions: Optional[Dict[str, PrecisionCondition]] = Field(default=None) + conditions: dict[str, PrecisionCondition] | None = Field(default=None) # This comes from the API - expected_end_date: Optional[AwareDatetimeISO] = Field( + expected_end_date: AwareDatetimeISO | None = Field( default=None, validation_alias="end_date" ) # This does not come from the API. We set it when we update this in the db. - created: Optional[AwareDatetimeISO] = Field(default=None) - updated: Optional[AwareDatetimeISO] = Field(default=None) + created: AwareDatetimeISO | None = Field(default=None) + updated: AwareDatetimeISO | None = Field(default=None) @property def internal_id(self) -> str: @@ -199,7 +198,7 @@ class PrecisionSurvey(MarketplaceTask): @computed_field @cached_property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set() for q in self.quotas: s.update(set(q.condition_hashes)) @@ -219,7 +218,7 @@ class PrecisionSurvey(MarketplaceTask): return data @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return PrecisionCondition @property @@ -227,7 +226,7 @@ class PrecisionSurvey(MarketplaceTask): return "age" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: PrecisionCondition( question_id="gender", @@ -262,7 +261,7 @@ class PrecisionSurvey(MarketplaceTask): exclude={"updated", "conditions", "created"} ) == other.model_dump(exclude={"updated", "conditions", "created"}) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump( mode="json", exclude={ @@ -283,11 +282,11 @@ class PrecisionSurvey(MarketplaceTask): return d @classmethod - def from_db(cls, d: Dict[str, Any]) -> Self: - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) + def from_db(cls, d: dict[str, Any]) -> Self: + d["created"] = d["created"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) d["expected_end_date"] = ( - d["expected_end_date"].replace(tzinfo=timezone.utc) + d["expected_end_date"].replace(tzinfo=UTC) if d["expected_end_date"] else None ) @@ -295,7 +294,7 @@ class PrecisionSurvey(MarketplaceTask): d["used_question_ids"] = json.loads(d["used_question_ids"]) return cls.model_validate(d) - def passes_quotas(self, criteria_evaluation: Dict[str, Optional[bool]]) -> bool: + def passes_quotas(self, criteria_evaluation: dict[str, bool | None]) -> bool: # We have to match 1 or more quota. # Quotas are exclusionary: they can NOT match a quota where currently_open=0 any_pass = False @@ -308,8 +307,8 @@ class PrecisionSurvey(MarketplaceTask): return any_pass def passes_quotas_soft( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], Set[str]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, set[str]]: # Quotas are exclusionary. They can NOT match a quota where currently_open=0 quota_eval = { quota: quota.matches_soft(criteria_evaluation) for quota in self.quotas @@ -345,19 +344,19 @@ class PrecisionSurvey(MarketplaceTask): return False, set() def determine_eligibility( - self, criteria_evaluation: Dict[str, Optional[bool]] + self, criteria_evaluation: dict[str, bool | None] ) -> bool: return self.is_open and self.passes_quotas(criteria_evaluation) def determine_eligibility_soft( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Tuple[Optional[bool], Optional[Set[str]]]: + self, criteria_evaluation: dict[str, bool | None] + ) -> tuple[bool | None, set[str] | None]: if not self.is_open: return False, None return self.passes_quotas_soft(criteria_evaluation) def participation_allowed( - self, att_survey_ids: Set[str], att_group_ids: Set[str] + self, att_survey_ids: set[str], att_group_ids: set[str] ) -> bool: """ Checks if this user can participate in this survey diff --git a/generalresearch/models/precision/task_collection.py b/generalresearch/models/precision/task_collection.py index 233d329..daea448 100644 --- a/generalresearch/models/precision/task_collection.py +++ b/generalresearch/models/precision/task_collection.py @@ -36,8 +36,8 @@ PrecisionTaskCollectionSchema = DataFrameSchema( "expected_end_date": Column(dtype=pd.DatetimeTZDtype(tz="UTC"), nullable=True), "created": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), "updated": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), - "used_question_ids": Column(List[str]), - "all_hashes": Column(List[str]), # set >> list for column support + "used_question_ids": Column(list[str]), + "all_hashes": Column(list[str]), # set >> list for column support }, checks=[], index=Index( @@ -53,10 +53,10 @@ PrecisionTaskCollectionSchema = DataFrameSchema( class PrecisionTaskCollection(TaskCollection): - items: List[PrecisionSurvey] + items: list[PrecisionSurvey] _schema = PrecisionTaskCollectionSchema - def to_row(self, s: PrecisionSurvey) -> Dict[str, Any]: + def to_row(self, s: PrecisionSurvey) -> dict[str, Any]: d = s.model_dump( mode="json", exclude={ diff --git a/generalresearch/models/prodege/__init__.py b/generalresearch/models/prodege/__init__.py index d419c0c..5c6659a 100644 --- a/generalresearch/models/prodege/__init__.py +++ b/generalresearch/models/prodege/__init__.py @@ -1,8 +1,7 @@ from enum import Enum -from typing import Literal +from typing import Annotated, Literal from pydantic import Field -from typing_extensions import Annotated ProdegeQuestionIdType = Annotated[ str, Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") diff --git a/generalresearch/models/prodege/question.py b/generalresearch/models/prodege/question.py index 1c61ab9..3ef4772 100644 --- a/generalresearch/models/prodege/question.py +++ b/generalresearch/models/prodege/question.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum from functools import cached_property from typing import TYPE_CHECKING, Any, Literal @@ -43,9 +43,7 @@ class ProdegeUserQuestionAnswer(BaseModel): # This may be a pipe-separated string if the question_type is multi. regex means any chars except capital letters option_id: str = Field(pattern=r"^[^A-Z]*$") - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # ISO 3166-1 alpha-2 (two-letter codes, lowercase) country_iso: str = Field( diff --git a/generalresearch/models/prodege/survey.py b/generalresearch/models/prodege/survey.py index c12f130..1601fa3 100644 --- a/generalresearch/models/prodege/survey.py +++ b/generalresearch/models/prodege/survey.py @@ -4,7 +4,7 @@ from __future__ import annotations import json import logging from collections import defaultdict -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from functools import cached_property from typing import Any, Literal, Type @@ -127,7 +127,7 @@ class ProdegeQuota(BaseModel): return self.remaining_count >= min_open_spots @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return ProdegeCondition @property @@ -276,7 +276,7 @@ class ProdegeUserPastParticipation(BaseModel): raise ValueError(f"Unknown ext_status_code_1: {self.ext_status_code_1}") def days_ago(self) -> float: - now = datetime.now(timezone.utc) + now = datetime.now(UTC) return (now - self.started).total_seconds() / (3600 * 24) @@ -486,7 +486,7 @@ class ProdegeSurvey(MarketplaceTask): return data @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return ProdegeCondition @property @@ -656,8 +656,8 @@ class ProdegeSurvey(MarketplaceTask): @classmethod def from_db(cls, d: dict[str, Any]) -> ProdegeSurvey: - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) + d["created"] = d["created"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) d["quotas"] = json.loads(d["quotas"]) for k in [ "max_clicks_settings", diff --git a/generalresearch/models/prodege/task_collection.py b/generalresearch/models/prodege/task_collection.py index 19e594f..d3e4a20 100644 --- a/generalresearch/models/prodege/task_collection.py +++ b/generalresearch/models/prodege/task_collection.py @@ -30,8 +30,8 @@ ProdegeTaskCollectionSchema = DataFrameSchema( "conversion_rate": Column(float, Check.between(0, 1), nullable=True), "created": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), "updated": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), - "used_question_ids": Column(List[str]), - "all_hashes": Column(List[str]), # set >> list for column support + "used_question_ids": Column(list[str]), + "all_hashes": Column(list[str]), # set >> list for column support "is_recontact": Column(bool), # Not including here: entrance_url, max_clicks_settings, past_participation, include_psids, exclude_psids, # quotas, source, conditions diff --git a/generalresearch/models/repdata/survey.py b/generalresearch/models/repdata/survey.py index 2290ca6..fa71c04 100644 --- a/generalresearch/models/repdata/survey.py +++ b/generalresearch/models/repdata/survey.py @@ -3,10 +3,10 @@ from __future__ import annotations import json import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from functools import cached_property -from typing import Any, Literal, Type +from typing import Any, Literal, Self, Type from uuid import UUID from pydantic import ( @@ -17,7 +17,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.grpc import timestamp_from_datetime from generalresearch.locales import Localelator @@ -304,7 +303,7 @@ class RepDataStream(MarketplaceTask): return self.stream_status == RepDataStatus.LIVE @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return RepDataCondition @property @@ -538,8 +537,8 @@ class RepDataSurveyHashed(RepDataSurvey): DeviceType(int(x)) for x in res["allowed_devices"].split(",") ] if res["created"] is not None: - res["created"] = res["created"].replace(tzinfo=timezone.utc) - res["last_updated"] = res["last_updated"].replace(tzinfo=timezone.utc) + res["created"] = res["created"].replace(tzinfo=UTC) + res["last_updated"] = res["last_updated"].replace(tzinfo=UTC) return cls.model_validate(res) def to_mysql(self) -> dict[str, Any]: @@ -553,7 +552,7 @@ class RepDataSurveyHashed(RepDataSurvey): return d def to_grpc(self, repdata_pb2): - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) timestamp = timestamp_from_datetime(now) return repdata_pb2.RepDataOpportunity( diff --git a/generalresearch/models/sago/__init__.py b/generalresearch/models/sago/__init__.py index 292f0f2..19e7b6d 100644 --- a/generalresearch/models/sago/__init__.py +++ b/generalresearch/models/sago/__init__.py @@ -1,7 +1,7 @@ from enum import Enum +from typing import Annotated from pydantic import Field -from typing_extensions import Annotated SagoQuestionIdType = Annotated[ str, Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") diff --git a/generalresearch/models/sago/survey.py b/generalresearch/models/sago/survey.py index e73ddee..0b5830b 100644 --- a/generalresearch/models/sago/survey.py +++ b/generalresearch/models/sago/survey.py @@ -2,14 +2,13 @@ from __future__ import annotations import json import logging -from datetime import timezone +from datetime import UTC, timezone from decimal import Decimal from functools import cached_property -from typing import Annotated, Any, Literal, Type +from typing import Annotated, Any, Literal, Self, Type from more_itertools import flatten from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import LogicalOperator, Source @@ -235,7 +234,7 @@ class SagoSurvey(MarketplaceTask): return data @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return SagoCondition @property @@ -314,9 +313,9 @@ class SagoSurvey(MarketplaceTask): @classmethod def from_db(cls, d: dict[str, Any]): - d["created"] = d["created"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) - d["modified_api"] = d["modified_api"].replace(tzinfo=timezone.utc) + d["created"] = d["created"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) + d["modified_api"] = d["modified_api"].replace(tzinfo=UTC) d["qualifications"] = json.loads(d["qualifications"]) d["used_question_ids"] = json.loads(d["used_question_ids"]) d["quotas"] = json.loads(d["quotas"]) diff --git a/generalresearch/models/spectrum/__init__.py b/generalresearch/models/spectrum/__init__.py index b62c089..0040551 100644 --- a/generalresearch/models/spectrum/__init__.py +++ b/generalresearch/models/spectrum/__init__.py @@ -1,7 +1,7 @@ from enum import Enum +from typing import Annotated from pydantic import Field -from typing_extensions import Annotated SpectrumQuestionIdType = Annotated[ str, Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") diff --git a/generalresearch/models/spectrum/question.py b/generalresearch/models/spectrum/question.py index db8a55d..81f8655 100644 --- a/generalresearch/models/spectrum/question.py +++ b/generalresearch/models/spectrum/question.py @@ -3,10 +3,10 @@ from __future__ import annotations import json import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Self from uuid import UUID from pydantic import ( @@ -16,7 +16,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.models import MAX_INT32, Source, string_utils from generalresearch.models.custom_types import AwareDatetimeISO @@ -50,9 +49,7 @@ class SpectrumUserQuestionAnswer(BaseModel): # This may be a pipe-separated string if the question_type is multi. regex # means any chars except capital letters option_id: str = Field(pattern=r"^[^A-Z]*$") - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # ISO 3166-1 alpha-2 (two-letter codes, lowercase) country_iso: str = Field( max_length=2, min_length=2, pattern=r"^[a-z]{2}$", frozen=True @@ -283,7 +280,7 @@ class SpectrumQuestion(MarketplaceQuestion): ] created = ( - datetime.utcfromtimestamp(d["crtd_on"] / 1000).replace(tzinfo=timezone.utc) + datetime.utcfromtimestamp(d["crtd_on"] / 1000).replace(tzinfo=UTC) if d.get("crtd_on") else None ) @@ -308,9 +305,7 @@ class SpectrumQuestion(MarketplaceQuestion): SpectrumQuestionOption(id=r["id"], text=r["text"], order=r["order"]) for r in d["options"] ] - d["created"] = ( - d["created"].replace(tzinfo=timezone.utc) if d["created"] else None - ) + d["created"] = d["created"].replace(tzinfo=UTC) if d["created"] else None return cls( question_id=d["question_id"], diff --git a/generalresearch/models/spectrum/survey.py b/generalresearch/models/spectrum/survey.py index a591445..f842f46 100644 --- a/generalresearch/models/spectrum/survey.py +++ b/generalresearch/models/spectrum/survey.py @@ -2,13 +2,12 @@ from __future__ import annotations import json import logging -from datetime import timezone +from datetime import UTC, timezone from decimal import Decimal -from typing import Any, Literal, Type +from typing import Any, Literal, Self, Type from more_itertools import flatten from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator -from typing_extensions import Self from generalresearch.locales import Localelator from generalresearch.models import Source, TaskCalculationType @@ -297,7 +296,7 @@ class SpectrumSurvey(MarketplaceTask): return data @property - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: return SpectrumCondition @property @@ -389,16 +388,14 @@ class SpectrumSurvey(MarketplaceTask): @classmethod def from_db(cls, d: dict[str, Any]) -> Self: - d["created_api"] = d["created_api"].replace(tzinfo=timezone.utc) - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) - d["modified_api"] = d["modified_api"].replace(tzinfo=timezone.utc) + d["created_api"] = d["created_api"].replace(tzinfo=UTC) + d["updated"] = d["updated"].replace(tzinfo=UTC) + d["modified_api"] = d["modified_api"].replace(tzinfo=UTC) d["field_end_date"] = ( - d["field_end_date"].replace(tzinfo=timezone.utc) - if d["field_end_date"] - else None + d["field_end_date"].replace(tzinfo=UTC) if d["field_end_date"] else None ) d["project_last_complete_date"] = ( - d["project_last_complete_date"].replace(tzinfo=timezone.utc) + d["project_last_complete_date"].replace(tzinfo=UTC) if d["project_last_complete_date"] else None ) diff --git a/generalresearch/models/string_utils.py b/generalresearch/models/string_utils.py index 23c1017..d76456f 100644 --- a/generalresearch/models/string_utils.py +++ b/generalresearch/models/string_utils.py @@ -2,7 +2,7 @@ import unicodedata from typing import Optional -def remove_nbsp(s: Optional[str]) -> Optional[str]: +def remove_nbsp(s: str | None) -> str | None: # Some text comes back from the API with lots of (copied from excel or # something), and random unicode... if s: diff --git a/generalresearch/models/thl/category.py b/generalresearch/models/thl/category.py index 1ed436a..ebfc840 100644 --- a/generalresearch/models/thl/category.py +++ b/generalresearch/models/thl/category.py @@ -1,10 +1,9 @@ from __future__ import annotations -from typing import Any +from typing import Any, Self from uuid import uuid4 from pydantic import BaseModel, Field, PositiveInt, model_validator -from typing_extensions import Self from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/models/thl/contest/__init__.py b/generalresearch/models/thl/contest/__init__.py index 363c8c0..0444586 100644 --- a/generalresearch/models/thl/contest/__init__.py +++ b/generalresearch/models/thl/contest/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone +from typing import Self from uuid import uuid4 from pydantic import ( @@ -10,7 +11,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Self from generalresearch.currency import USDCent from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr @@ -106,7 +106,7 @@ class ContestWinner(BaseModel): uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) created_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When this user won this prize", ) diff --git a/generalresearch/models/thl/contest/contest.py b/generalresearch/models/thl/contest/contest.py index 6889dcb..8c45f13 100644 --- a/generalresearch/models/thl/contest/contest.py +++ b/generalresearch/models/thl/contest/contest.py @@ -2,8 +2,8 @@ from __future__ import annotations import json from abc import ABC, abstractmethod -from datetime import datetime, timezone -from typing import Any +from datetime import UTC, datetime, timezone +from typing import Any, Self from uuid import uuid4 from pydantic import ( @@ -14,7 +14,6 @@ from pydantic import ( NonNegativeInt, model_validator, ) -from typing_extensions import Self from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.thl.contest import ( @@ -57,7 +56,7 @@ class ContestBase(BaseModel, ABC): starts_at: AwareDatetimeISO = Field( description="When the contest starts", - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), ) terms_and_conditions: HttpUrl | None = Field(default=None) @@ -91,11 +90,11 @@ class Contest(ContestBase): product_id: UUIDStr = Field(description="Contest applies only to a single BP") created_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When this contest was created", ) updated_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When this contest was last modified. Does not include " "entries being created/modified", ) @@ -139,7 +138,7 @@ class Contest(ContestBase): def should_end(self) -> tuple[bool, ContestEndReason | None]: if self.status == ContestStatus.ACTIVE: if self.end_condition.ends_at: - if datetime.now(tz=timezone.utc) >= self.end_condition.ends_at: + if datetime.now(tz=UTC) >= self.end_condition.ends_at: return True, ContestEndReason.ENDS_AT return False, None @@ -158,14 +157,14 @@ class Contest(ContestBase): if winners is not None: self.update( status=ContestStatus.COMPLETED, - ended_at=datetime.now(tz=timezone.utc), + ended_at=datetime.now(tz=UTC), end_reason=reason, all_winners=winners, ) else: self.update( status=ContestStatus.COMPLETED, - ended_at=datetime.now(tz=timezone.utc), + ended_at=datetime.now(tz=UTC), end_reason=reason, ) return None @@ -211,7 +210,7 @@ class ContestUserView(Contest): ) def is_user_eligible(self, country_iso: str) -> tuple[bool, str]: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) assert country_iso.lower() == country_iso if now < self.starts_at: diff --git a/generalresearch/models/thl/contest/contest_entry.py b/generalresearch/models/thl/contest/contest_entry.py index cddae14..31ef317 100644 --- a/generalresearch/models/thl/contest/contest_entry.py +++ b/generalresearch/models/thl/contest/contest_entry.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from uuid import uuid4 from pydantic import ( @@ -39,12 +39,8 @@ class ContestEntryCreate(BaseModel): class ContestEntry(BaseModel): uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) - created_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(timezone.utc) - ) - updated_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(timezone.utc) - ) + created_at: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(UTC)) + updated_at: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(UTC)) # entry_type and amount are the same as on ContestEntryCreate entry_type: ContestEntryType = Field() diff --git a/generalresearch/models/thl/contest/io.py b/generalresearch/models/thl/contest/io.py index e68f76e..c6af719 100644 --- a/generalresearch/models/thl/contest/io.py +++ b/generalresearch/models/thl/contest/io.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from uuid import uuid4 from generalresearch.models.thl.contest.definitions import ContestType @@ -37,7 +37,7 @@ from generalresearch.models.thl.contest.contest import Contest def contest_create_to_contest( product_id: str, contest_create: ContestCreate ) -> Contest: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) d = contest_create.model_dump(mode="json") d["uuid"] = uuid4().hex d["product_id"] = product_id diff --git a/generalresearch/models/thl/contest/leaderboard.py b/generalresearch/models/thl/contest/leaderboard.py index 8167f46..c5e0626 100644 --- a/generalresearch/models/thl/contest/leaderboard.py +++ b/generalresearch/models/thl/contest/leaderboard.py @@ -1,7 +1,7 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone -from typing import Any, Literal +from datetime import UTC, datetime, timedelta, timezone +from typing import Any, Literal, Self from pydantic import ( ConfigDict, @@ -11,7 +11,6 @@ from pydantic import ( model_validator, ) from redis import Redis -from typing_extensions import Self from generalresearch.decorators import LOG from generalresearch.managers.leaderboard import country_timezone @@ -195,7 +194,7 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): def should_end(self) -> tuple[bool, ContestEndReason | None]: if self.status == ContestStatus.ACTIVE: if self.end_condition.ends_at: - if datetime.now(tz=timezone.utc) >= self.end_condition.ends_at: + if datetime.now(tz=UTC) >= self.end_condition.ends_at: return True, ContestEndReason.ENDS_AT return False, None @@ -276,7 +275,7 @@ class LeaderboardContestUserView(LeaderboardContest, ContestUserView): if self.user_winnings: return False, "User already won" - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) if self.leaderboard_model.period_end_utc < now: return False, "Contest is over" if self.leaderboard_model.period_start_utc > now: diff --git a/generalresearch/models/thl/contest/milestone.py b/generalresearch/models/thl/contest/milestone.py index f62be8f..8b74d50 100644 --- a/generalresearch/models/thl/contest/milestone.py +++ b/generalresearch/models/thl/contest/milestone.py @@ -10,7 +10,7 @@ from pydantic import ( Field, PositiveInt, ) -from typing_extensions import Self +from typing import Self from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.thl.contest.contest import ( diff --git a/generalresearch/models/thl/contest/raffle.py b/generalresearch/models/thl/contest/raffle.py index d3157e3..b497a44 100644 --- a/generalresearch/models/thl/contest/raffle.py +++ b/generalresearch/models/thl/contest/raffle.py @@ -3,8 +3,8 @@ from __future__ import annotations import logging import random from collections import defaultdict -from datetime import datetime, timezone -from typing import Any, Literal +from datetime import UTC, datetime, timezone +from typing import Any, Literal, Self from pydantic import ( ConfigDict, @@ -14,7 +14,6 @@ from pydantic import ( model_validator, ) from scipy.stats import hypergeom -from typing_extensions import Self from generalresearch.currency import USDCent from generalresearch.models.thl.contest import ( @@ -202,7 +201,7 @@ class RaffleContest(RaffleContestCreate, Contest): c = self.end_condition if c.target_entry_amount and self.current_amount >= c.target_entry_amount: return True - if c.ends_at and datetime.now(tz=timezone.utc) >= c.ends_at: + if c.ends_at and datetime.now(tz=UTC) >= c.ends_at: return True return False diff --git a/generalresearch/models/thl/finance.py b/generalresearch/models/thl/finance.py index b72ecf6..a992f78 100644 --- a/generalresearch/models/thl/finance.py +++ b/generalresearch/models/thl/finance.py @@ -1,7 +1,7 @@ from __future__ import annotations import random -from datetime import timezone +from datetime import UTC, timezone from typing import TYPE_CHECKING from uuid import uuid4 @@ -16,8 +16,8 @@ from pydantic import ( model_validator, ) from pydantic.json_schema import SkipJsonSchema -from generalresearch.config import is_debug +from generalresearch.config import is_debug from generalresearch.currency import USDCent from generalresearch.decorators import LOG from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr @@ -150,7 +150,7 @@ class POPFinancial(BaseModel): index: int # Not useful, just a RangeIndex row: pd.DataFrame - row["time_idx"] = row.time_idx.to_pydatetime().replace(tzinfo=timezone.utc) + row["time_idx"] = row.time_idx.to_pydatetime().replace(tzinfo=UTC) instance = ProductBalances.from_pandas(row) res.append( @@ -842,12 +842,12 @@ class BusinessBalances(BaseModel): from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) + from generalresearch.managers.thl.product import ProductManager from generalresearch.models.thl.ledger import ( AccountType, Direction, ) from generalresearch.models.thl.product import Product - from generalresearch.managers.thl.product import ProductManager # Validate the input accounts assert len(accounts) > 0, "Must provide accounts" diff --git a/generalresearch/models/thl/ipinfo.py b/generalresearch/models/thl/ipinfo.py index 0d254e2..3f212cf 100644 --- a/generalresearch/models/thl/ipinfo.py +++ b/generalresearch/models/thl/ipinfo.py @@ -1,8 +1,8 @@ from __future__ import annotations import ipaddress -from datetime import datetime, timezone -from typing import Any, Literal +from datetime import UTC, datetime, timezone +from typing import Any, Literal, Self from faker import Faker from pydantic import ( @@ -13,7 +13,6 @@ from pydantic import ( PrivateAttr, field_validator, ) -from typing_extensions import Self from generalresearch.models.custom_types import ( AwareDatetimeISO, @@ -95,7 +94,7 @@ class IPGeoname(BaseModel): is_in_european_union: bool | None = Field(default=None) updated: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), ) @field_validator( @@ -119,7 +118,7 @@ class IPGeoname(BaseModel): @classmethod def from_mysql(cls, d: dict[str, Any]) -> Self: - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) + d["updated"] = d["updated"].replace(tzinfo=UTC) return cls.model_validate(d) @@ -205,7 +204,7 @@ class IPInformation(BaseModel): ) updated: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), ) _geoname: IPGeoname | None = PrivateAttr(default=None) @@ -255,7 +254,7 @@ class IPInformation(BaseModel): @classmethod def from_mysql(cls, d: dict) -> Self: - d["updated"] = d["updated"].replace(tzinfo=timezone.utc) + d["updated"] = d["updated"].replace(tzinfo=UTC) return cls.model_validate(d) diff --git a/generalresearch/models/thl/leaderboard.py b/generalresearch/models/thl/leaderboard.py index 399a906..dce3280 100644 --- a/generalresearch/models/thl/leaderboard.py +++ b/generalresearch/models/thl/leaderboard.py @@ -2,10 +2,11 @@ from __future__ import annotations import logging import math -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from enum import Enum from typing import Literal from uuid import UUID, uuid3 +from zoneinfo import ZoneInfo import pandas as pd from pydantic import ( @@ -17,7 +18,6 @@ from pydantic import ( field_validator, model_validator, ) -from zoneinfo import ZoneInfo from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.legacy.api_status import StatusResponse @@ -167,13 +167,13 @@ class Leaderboard(BaseModel): def period_start_utc(self) -> datetime: # The start of the time period covered by this board in UTC, tz-aware # e.g. datetime(2024, 7, 12, 4, 0, 0, 0, tzinfo=timezone.utc) - return self.period_start_local.astimezone(timezone.utc) + return self.period_start_local.astimezone(UTC) @property def period_end_utc(self) -> datetime: # The end of the time period covered by this board in UTC, tz-aware # e.g. datetime(2024, 7, 13, 3, 59, 59, 999999, tzinfo=timezone.utc) - return self.period_end_local.astimezone(timezone.utc) + return self.period_end_local.astimezone(UTC) @computed_field( description="(unix timestamp) The start time of the time range this leaderboard covers.", diff --git a/generalresearch/models/thl/ledger.py b/generalresearch/models/thl/ledger.py index 3f8b123..dd37d98 100644 --- a/generalresearch/models/thl/ledger.py +++ b/generalresearch/models/thl/ledger.py @@ -1,8 +1,8 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum -from typing import Annotated, Any, Literal, Union +from typing import Annotated, Any, Literal, Self, Union from uuid import uuid4 from pydantic import ( @@ -15,7 +15,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.models.custom_types import ( AwareDatetimeISO, @@ -253,7 +252,7 @@ class LedgerTransaction(BaseModel): id: int | None = Field(default=None) created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When the Transaction (TX) was created into the database." "This does not represent the exact time for any action" "which may be responsible for this Transaction (TX), and " @@ -283,9 +282,7 @@ class LedgerTransaction(BaseModel): """Created should not be in the future. This will mess up LedgerAccountStatement / groupby rollups. """ - assert ( - datetime.now(tz=timezone.utc) > created - ), "created cannot be in the future" + assert datetime.now(tz=UTC) > created, "created cannot be in the future" return created @field_validator("entries", mode="after") @@ -536,12 +533,10 @@ class UserLedgerTransactionTaskAdjustment(UserLedgerTransaction): UserLedgerTransactionType = Annotated[ - Union[ - UserLedgerTransactionUserPayout, - UserLedgerTransactionUserBonus, - UserLedgerTransactionTaskAdjustment, - UserLedgerTransactionTaskComplete, - ], + UserLedgerTransactionUserPayout + | UserLedgerTransactionUserBonus + | UserLedgerTransactionTaskAdjustment + | UserLedgerTransactionTaskComplete, Field(discriminator="tx_type"), ] diff --git a/generalresearch/models/thl/ledger_example.py b/generalresearch/models/thl/ledger_example.py index 767be85..92ad83d 100644 --- a/generalresearch/models/thl/ledger_example.py +++ b/generalresearch/models/thl/ledger_example.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import Any from uuid import uuid4 @@ -16,7 +16,7 @@ def _example_user_tx_payout(schema: dict[str, Any]) -> None: amount=-5, description="HIT Reward", payout_format="${payout/100:.2f}", - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ).model_dump(mode="json") @@ -30,7 +30,7 @@ def _example_user_tx_bonus(schema: dict[str, Any]) -> None: amount=100, description="Compensation Bonus", payout_format="${payout/100:.2f}", - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ).model_dump(mode="json") @@ -44,7 +44,7 @@ def _example_user_tx_complete(schema: dict[str, Any]) -> None: amount=38, description="Task Complete", payout_format="${payout/100:.2f}", - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), tsid=uuid4().hex, ).model_dump(mode="json") @@ -59,6 +59,6 @@ def _example_user_tx_adjustment(schema: dict[str, Any]) -> None: amount=-38, description="Task Adjustment", payout_format="${payout/100:.2f}", - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), tsid=uuid4().hex, ).model_dump(mode="json") diff --git a/generalresearch/models/thl/offerwall/__init__.py b/generalresearch/models/thl/offerwall/__init__.py index d2d7d36..3acc592 100644 --- a/generalresearch/models/thl/offerwall/__init__.py +++ b/generalresearch/models/thl/offerwall/__init__.py @@ -4,7 +4,7 @@ import hashlib import json from decimal import Decimal from enum import Enum -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import ( BaseModel, @@ -13,7 +13,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Self from generalresearch.models import Source from generalresearch.models.custom_types import IPvAnyAddressStr diff --git a/generalresearch/models/thl/offerwall/base.py b/generalresearch/models/thl/offerwall/base.py index 3a867b6..33489df 100644 --- a/generalresearch/models/thl/offerwall/base.py +++ b/generalresearch/models/thl/offerwall/base.py @@ -4,7 +4,7 @@ import statistics from datetime import timedelta from decimal import Decimal from string import Formatter -from typing import Any +from typing import Annotated, Any, Self from uuid import uuid4 import numpy as np @@ -18,7 +18,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Annotated, Self from generalresearch.models import Source from generalresearch.models.custom_types import HttpsUrl, UUIDStr diff --git a/generalresearch/models/thl/offerwall/cache.py b/generalresearch/models/thl/offerwall/cache.py index c36568e..82ab36d 100644 --- a/generalresearch/models/thl/offerwall/cache.py +++ b/generalresearch/models/thl/offerwall/cache.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import Any from pydantic import BaseModel, Field @@ -26,9 +26,7 @@ class GetOfferWallCache(BaseModel): request_id: str = Field() offerwall: OfferwallBase = Field() all_sids: list[str] = Field() - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(UTC)) latest_ip_info: dict[str, Any] = Field( description="So we can easily check if user's IP info has changed" ) @@ -51,9 +49,7 @@ class SessionInfoCache(BaseModel): # will get pruned as tasks are attempted tasks: list[ScoredTaskResult] = Field() - started: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + started: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # The count of attempts per marketplace mp_retry_count: dict[Source, int] = Field(default_factory=dict) diff --git a/generalresearch/models/thl/payout_format.py b/generalresearch/models/thl/payout_format.py index 9f22ace..a7fc4fb 100644 --- a/generalresearch/models/thl/payout_format.py +++ b/generalresearch/models/thl/payout_format.py @@ -2,9 +2,9 @@ from __future__ import annotations import decimal import re +from typing import Annotated from pydantic import AfterValidator, Field -from typing_extensions import Annotated # Matches only digits, parenthesis, + , -, *, / and the string payout. xform_format_re = re.compile(pattern=r"^[\d()+\-*/.]*payout[\d()+\-*/.]*$") diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index 9b7d66a..8377e45 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -6,14 +6,15 @@ import json import math import warnings from collections import defaultdict +from collections.abc import Callable from decimal import Decimal from enum import Enum from functools import cached_property, partial from typing import ( TYPE_CHECKING, Any, - Callable, Literal, + Self, ) from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit from uuid import uuid4 @@ -34,7 +35,6 @@ from pydantic import ( model_validator, ) from pydantic.json_schema import SkipJsonSchema -from typing_extensions import Self from generalresearch.currency import USDCent from generalresearch.decorators import LOG @@ -941,9 +941,7 @@ class Product(BaseModel, validate_assignment=True): # Initialization is deferred until unless it's called # (see .prebuild_***()) - balance: ProductBalances | None = Field( - default=None, description="Product Balance" - ) + balance: ProductBalances | None = Field(default=None, description="Product Balance") payouts_total_str: str | None = Field(default=None) payouts_total: USDCent | None = Field(default=None) diff --git a/generalresearch/models/thl/profiling/marketplace.py b/generalresearch/models/thl/profiling/marketplace.py index 027aa4c..9038cf6 100644 --- a/generalresearch/models/thl/profiling/marketplace.py +++ b/generalresearch/models/thl/profiling/marketplace.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from functools import cached_property from typing import Any @@ -111,9 +111,7 @@ class MarketplaceUserQuestionAnswer(BaseModel): # This may be a pipe-separated string if the question_type is multi. Regex # means any chars except capital letters option_id: str = Field(pattern=r"^[^A-Z]*$") - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) country_iso: CountryISO = Field(frozen=True) language_iso: LanguageISO = Field(frozen=True) diff --git a/generalresearch/models/thl/profiling/other_option.py b/generalresearch/models/thl/profiling/other_option.py index ae58e90..6d789e5 100644 --- a/generalresearch/models/thl/profiling/other_option.py +++ b/generalresearch/models/thl/profiling/other_option.py @@ -40,7 +40,7 @@ texts_in = { } -def option_is_catch_all(c: "UpkQuestionChoice") -> bool: +def option_is_catch_all(c: UpkQuestionChoice) -> bool: """ Exclusive not specifically in the sense that it is a multi-select question and if this option is selected no others can be selected. But also in the diff --git a/generalresearch/models/thl/profiling/upk_question.py b/generalresearch/models/thl/profiling/upk_question.py index 307bc33..78f9511 100644 --- a/generalresearch/models/thl/profiling/upk_question.py +++ b/generalresearch/models/thl/profiling/upk_question.py @@ -5,7 +5,7 @@ import json import re from enum import Enum from functools import cached_property -from typing import Any, List, Literal, Union +from typing import Annotated, Any, List, Literal, Union from pydantic import ( BaseModel, @@ -16,7 +16,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Annotated from generalresearch.models import Source from generalresearch.models.custom_types import UUIDStr @@ -215,11 +214,9 @@ SelectorType = ( | UpkQuestionSelectorHIDDEN ) Configuration = Annotated[ - Union[ - UpkQuestionConfigurationMC, - UpkQuestionConfigurationTE, - UpkQuestionConfigurationSLIDER, - ], + UpkQuestionConfigurationMC + | UpkQuestionConfigurationTE + | UpkQuestionConfigurationSLIDER, Field(discriminator="type"), ] @@ -433,7 +430,7 @@ class UpkQuestion(BaseModel): @field_validator("choices") @classmethod - def order_choices(cls, choices: List): + def order_choices(cls, choices: list): if choices: choices.sort(key=lambda x: x.order) return choices diff --git a/generalresearch/models/thl/profiling/upk_question_answer.py b/generalresearch/models/thl/profiling/upk_question_answer.py index 0024e68..c59d99d 100644 --- a/generalresearch/models/thl/profiling/upk_question_answer.py +++ b/generalresearch/models/thl/profiling/upk_question_answer.py @@ -1,7 +1,7 @@ from __future__ import annotations -from datetime import datetime, timezone -from typing import Any +from datetime import UTC, datetime, timezone +from typing import Any, Self from uuid import uuid4 from pydantic import ( @@ -12,7 +12,6 @@ from pydantic import ( computed_field, model_validator, ) -from typing_extensions import Self from generalresearch.models import MAX_INT32 from generalresearch.models.custom_types import ( @@ -60,9 +59,7 @@ class UpkQuestionAnswer(BaseModel): # ISO 3166-1 alpha-2 (two-letter codes, lowercase) country_iso: CountryISOLike = Field() - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # If the property is PropertyType.UPK_ITEM, it should have an item (and no value). # If the property is UPK_NUMERICAL or UPK_TEXT, it'll have a value (and no item). diff --git a/generalresearch/models/thl/profiling/user_question_answer.py b/generalresearch/models/thl/profiling/user_question_answer.py index 8248623..a55b205 100644 --- a/generalresearch/models/thl/profiling/user_question_answer.py +++ b/generalresearch/models/thl/profiling/user_question_answer.py @@ -1,8 +1,9 @@ from __future__ import annotations import json -from datetime import datetime, timedelta, timezone -from typing import Any, Iterator, Literal +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta, timezone +from typing import Any, Literal, Self from pydantic import ( BaseModel, @@ -12,7 +13,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.grpc import timestamp_to_datetime from generalresearch.models import MAX_INT32, Source @@ -28,9 +28,7 @@ class UserQuestionAnswer(BaseModel): user_id: PositiveInt | None = Field(lt=MAX_INT32, default=None) question_id: UUIDStr = Field() answer: tuple[str, ...] = Field() - timestamp: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + timestamp: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) country_iso: CountryISO | Literal["xx"] = Field() language_iso: LanguageISO | Literal["xxx"] = Field() @@ -105,7 +103,7 @@ class UserQuestionAnswer(BaseModel): return True, "" def is_stale(self) -> bool: - return self.timestamp < datetime.now(tz=timezone.utc) - timedelta(days=30) + return self.timestamp < datetime.now(tz=UTC) - timedelta(days=30) @classmethod def from_grpc(cls, msg, default_timestamp: datetime) -> Self: @@ -129,7 +127,7 @@ class UserQuestionAnswer(BaseModel): DUMMY_UQA = UserQuestionAnswer( question_id="f118edd01cf1476ba7200a175fb4351d", answer=("0",), - timestamp=datetime(2020, 1, 1, tzinfo=timezone.utc), + timestamp=datetime(2020, 1, 1, tzinfo=UTC), country_iso="xx", language_iso="xxx", property_code="dummy", diff --git a/generalresearch/models/thl/session.py b/generalresearch/models/thl/session.py index 17142f3..c8e681c 100644 --- a/generalresearch/models/thl/session.py +++ b/generalresearch/models/thl/session.py @@ -2,9 +2,9 @@ from __future__ import annotations import json import logging -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, Self from uuid import uuid4 from pydantic import ( @@ -17,7 +17,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.models import DeviceType, Source from generalresearch.models.custom_types import ( @@ -69,9 +68,7 @@ class WallBase(BaseModel): buyer_id: str | None = Field(default=None, max_length=32) req_survey_id: str = Field(max_length=32) req_cpi: Decimal = Field(decimal_places=5, lt=1000, ge=0) - started: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + started: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # These get set on creation, or updated when the wall event is finished. So # they shouldn't really ever be NULL, but you don't have to pass them in @@ -158,9 +155,7 @@ class WallBase(BaseModel): @model_validator(mode="after") def check_timestamps(self): - assert self.started <= datetime.now( - tz=timezone.utc - ), "Started must not be in the future" + assert self.started <= datetime.now(tz=UTC), "Started must not be in the future" if self.finished: assert self.finished > self.started, "Finished must be after started" assert self.finished - self.started <= timedelta( @@ -278,7 +273,7 @@ class WallBase(BaseModel): # This is just used in tests at the moment. This needs to be adjusted. if finished is None: - finished = datetime.now(tz=timezone.utc) + finished = datetime.now(tz=UTC) self.update( status=status, @@ -313,7 +308,7 @@ class WallBase(BaseModel): ext_status_code_3, ) if finished is None: - finished = datetime.now(tz=timezone.utc) + finished = datetime.now(tz=UTC) self.update( status=status, status_code_1=status_code_1, @@ -386,7 +381,7 @@ class WallBase(BaseModel): TODO: Transition this over to use the ReportTask pydantic model. """ report_timestamp = ( - report_timestamp if report_timestamp else datetime.now(tz=timezone.utc) + report_timestamp if report_timestamp else datetime.now(tz=UTC) ) if self.status is None and self.finished is None: self.status = Status.ABANDON @@ -587,9 +582,7 @@ class Session(BaseModel): id: int | None = None uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) user: User - started: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + started: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # This is the "bucket" the user clicked on to start this session. We only # store the 4 fields: loi_min, loi_max, user_payout_min, user_payout_max @@ -881,7 +874,7 @@ class Session(BaseModel): if ( last_wall.status is None and self.status is None - and datetime.now(tz=timezone.utc) + and datetime.now(tz=UTC) > self.started + timedelta(seconds=task_timeout_seconds) ): last_wall.status = Status.TIMEOUT @@ -962,7 +955,7 @@ class Session(BaseModel): self, max_session_len: timedelta, max_session_hard_retry: int ) -> bool: - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) last_wall = self.get_last_visible_wall() if last_wall and last_wall.status == Status.COMPLETE: diff --git a/generalresearch/models/thl/survey/__init__.py b/generalresearch/models/thl/survey/__init__.py index b6ac740..d6ba895 100644 --- a/generalresearch/models/thl/survey/__init__.py +++ b/generalresearch/models/thl/survey/__init__.py @@ -108,7 +108,7 @@ class MarketplaceTask(BaseModel, ABC): @property @abstractmethod - def condition_model(self) -> Type[MarketplaceCondition]: + def condition_model(self) -> type[MarketplaceCondition]: """ The Condition Model for this survey class """ diff --git a/generalresearch/models/thl/survey/buyer.py b/generalresearch/models/thl/survey/buyer.py index 6d4d7a1..384bab4 100644 --- a/generalresearch/models/thl/survey/buyer.py +++ b/generalresearch/models/thl/survey/buyer.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from math import log from typing import Annotated @@ -46,7 +46,7 @@ class Buyer(BaseModel): ) label: str | None = Field(default=None, max_length=255) created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When this entry was made, or when the buyer was first seen", ) diff --git a/generalresearch/models/thl/survey/condition.py b/generalresearch/models/thl/survey/condition.py index 927b7e1..a85073c 100644 --- a/generalresearch/models/thl/survey/condition.py +++ b/generalresearch/models/thl/survey/condition.py @@ -4,7 +4,7 @@ import hashlib from abc import ABC from enum import Enum from functools import cached_property -from typing import Any +from typing import Annotated, Any, Self from pydantic import ( BaseModel, @@ -16,7 +16,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Annotated, Self from generalresearch.models import LogicalOperator diff --git a/generalresearch/models/thl/survey/model.py b/generalresearch/models/thl/survey/model.py index 3794c00..57bcbe2 100644 --- a/generalresearch/models/thl/survey/model.py +++ b/generalresearch/models/thl/survey/model.py @@ -1,8 +1,8 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal -from typing import Any +from typing import Annotated, Any from pydantic import ( BaseModel, @@ -15,7 +15,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Annotated from generalresearch.managers.thl.buyer import Buyer from generalresearch.models import Source @@ -71,12 +70,8 @@ class Survey(BaseModel): min_length=1, max_length=128, default=None, examples=["124"] ) - created_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) - updated_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created_at: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) + updated_at: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) is_live: bool = Field(default=True) is_recontact: bool = Field(default=False) @@ -188,9 +183,7 @@ class SurveyStat(BaseModel): # ---- Metadata ---- - updated_at: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + updated_at: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) @property def natural_key(self) -> str: diff --git a/generalresearch/models/thl/survey/penalty.py b/generalresearch/models/thl/survey/penalty.py index e9515d4..05153fe 100644 --- a/generalresearch/models/thl/survey/penalty.py +++ b/generalresearch/models/thl/survey/penalty.py @@ -1,11 +1,10 @@ from __future__ import annotations import abc -from datetime import datetime, timezone -from typing import Literal +from datetime import UTC, datetime, timezone +from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, TypeAdapter -from typing_extensions import Annotated from generalresearch.models import Source from generalresearch.models.custom_types import ( @@ -29,9 +28,7 @@ class SurveyPenalty(BaseModel, abc.ABC): penalty: float = Field(ge=0, le=1) - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) @property def sid(self): diff --git a/generalresearch/models/thl/task_adjustment.py b/generalresearch/models/thl/task_adjustment.py index 89a3873..1834898 100644 --- a/generalresearch/models/thl/task_adjustment.py +++ b/generalresearch/models/thl/task_adjustment.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from uuid import uuid4 @@ -26,11 +26,11 @@ class TaskAdjustmentEvent(BaseModel): uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When this event was created in the db", ) alerted: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), description="When we were notified about this change", ) diff --git a/generalresearch/models/thl/task_status.py b/generalresearch/models/thl/task_status.py index ee713a5..011d743 100644 --- a/generalresearch/models/thl/task_status.py +++ b/generalresearch/models/thl/task_status.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import datetime -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, Self from pydantic import ( BaseModel, @@ -12,7 +12,6 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self from generalresearch.models.custom_types import ( AwareDatetimeISO, diff --git a/generalresearch/models/thl/user.py b/generalresearch/models/thl/user.py index 355e331..55bbd18 100644 --- a/generalresearch/models/thl/user.py +++ b/generalresearch/models/thl/user.py @@ -3,8 +3,8 @@ from __future__ import annotations import json import logging import re -from datetime import datetime, timezone -from typing import TYPE_CHECKING +from datetime import UTC, datetime, timezone +from typing import TYPE_CHECKING, Annotated, Self from uuid import UUID, uuid4 from pydantic import ( @@ -19,7 +19,6 @@ from pydantic import ( model_validator, ) from sentry_sdk import set_tag, set_user -from typing_extensions import Annotated, Self from generalresearch.models import MAX_INT32 from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr @@ -126,7 +125,7 @@ class User(BaseModel): def check_not_in_future(cls, v: AwareDatetime) -> AwareDatetime: if v is not None: try: - assert v < datetime.now(tz=timezone.utc) + assert v < datetime.now(tz=UTC) except Exception: raise ValueError("Input is in the future") return v @@ -137,7 +136,7 @@ class User(BaseModel): def check_after_anno_domini(cls, v: AwareDatetime) -> AwareDatetime: if v is not None: try: - assert v > datetime(year=2016, month=7, day=13, tzinfo=timezone.utc) + assert v > datetime(year=2016, month=7, day=13, tzinfo=UTC) except Exception: raise ValueError("Input is before Anno Domini") return v @@ -294,9 +293,9 @@ class User(BaseModel): @classmethod def from_db(cls, res) -> Self: if res["created"]: - res["created"] = res["created"].replace(tzinfo=timezone.utc) + res["created"] = res["created"].replace(tzinfo=UTC) if res["last_seen"]: - res["last_seen"] = res["last_seen"].replace(tzinfo=timezone.utc) + res["last_seen"] = res["last_seen"].replace(tzinfo=UTC) res["product_id"] = UUID(res["product_id"]).hex res["uuid"] = UUID(res["uuid"]).hex return cls( diff --git a/generalresearch/models/thl/user_iphistory.py b/generalresearch/models/thl/user_iphistory.py index 5892d41..469f8ba 100644 --- a/generalresearch/models/thl/user_iphistory.py +++ b/generalresearch/models/thl/user_iphistory.py @@ -1,7 +1,8 @@ from __future__ import annotations import ipaddress -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone +from typing import Self from faker import Faker from pydantic import ( @@ -11,7 +12,6 @@ from pydantic import ( PositiveInt, field_validator, ) -from typing_extensions import Self from generalresearch.models.custom_types import ( AwareDatetimeISO, @@ -113,7 +113,7 @@ class IPRecord(BaseModel): # --- ORM --- @classmethod def from_mysql(cls, d: dict) -> Self: - created = d["created"].replace(tzinfo=timezone.utc) + created = d["created"].replace(tzinfo=UTC) d["created"] = created d["forwarded_ip_records"] = [] @@ -169,7 +169,7 @@ class UserIPHistory(BaseModel): def ips_timestamp(cls, ips): if ips is None: return None - cutoff = datetime.now(tz=timezone.utc) - timedelta(days=28) + cutoff = datetime.now(tz=UTC) - timedelta(days=28) return sorted( [x for x in ips if x.created > cutoff], key=lambda x: x.created, diff --git a/generalresearch/models/thl/user_profile.py b/generalresearch/models/thl/user_profile.py index e96266a..0ec605a 100644 --- a/generalresearch/models/thl/user_profile.py +++ b/generalresearch/models/thl/user_profile.py @@ -1,7 +1,7 @@ from __future__ import annotations import hashlib -from typing import Any +from typing import Annotated, Any, Self from pydantic import ( BaseModel, @@ -12,7 +12,6 @@ from pydantic import ( computed_field, ) from pydantic.json_schema import SkipJsonSchema -from typing_extensions import Annotated, Self from generalresearch.models import MAX_INT32, Source from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/models/thl/user_quality_event.py b/generalresearch/models/thl/user_quality_event.py index d6ebddc..52903e3 100644 --- a/generalresearch/models/thl/user_quality_event.py +++ b/generalresearch/models/thl/user_quality_event.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from enum import Enum from typing import Literal @@ -61,9 +61,7 @@ class TaskAdjustmentEvent(BaseModel): mid: UUIDStr = Field() source: Source = Field() status: WallAdjustedStatus = Field() - alert_time: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + alert_time: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) quality_event_type: Literal[QualityEventType.task_adjustment] = Field( default=QualityEventType.task_adjustment ) diff --git a/generalresearch/models/thl/userhealth.py b/generalresearch/models/thl/userhealth.py index e556dc8..fb15572 100644 --- a/generalresearch/models/thl/userhealth.py +++ b/generalresearch/models/thl/userhealth.py @@ -1,11 +1,10 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from enum import Enum -from typing import Dict, Optional +from typing import Dict, Optional, Self from pydantic import BaseModel, Field, NonNegativeFloat, PositiveInt -from typing_extensions import Self from generalresearch.models.custom_types import AwareDatetimeISO @@ -26,12 +25,12 @@ class AuditLog(BaseModel): are related to a User """ - id: Optional[PositiveInt] = Field(default=None) + id: PositiveInt | None = Field(default=None) user_id: PositiveInt = Field() created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc), - examples=[datetime.now(tz=timezone.utc)], + default_factory=lambda: datetime.now(tz=UTC), + examples=[datetime.now(tz=UTC)], description="When did this event occur", ) @@ -51,14 +50,14 @@ class AuditLog(BaseModel): # e.g. "upk-audit", "ip-audit", "entrance-limit" event_type: str = Field(max_length=64, examples=["entrance-limit"]) - event_msg: Optional[str] = Field( + event_msg: str | None = Field( default=None, min_length=3, max_length=256, description="The event message. Could be displayed on user's page", ) - event_value: Optional[NonNegativeFloat] = Field( + event_value: NonNegativeFloat | None = Field( default=None, description="Optionally store a numeric value associated with this " "event. For e.g. if we recalculate the user's normalized " @@ -68,12 +67,12 @@ class AuditLog(BaseModel): examples=[0.42], ) - def model_dump_mysql(self, **kwargs) -> Dict: + def model_dump_mysql(self, **kwargs) -> dict: d = self.model_dump(mode="json", **kwargs) d["created"] = self.created.replace(tzinfo=None) return d @classmethod - def from_mysql(cls, d: Dict) -> Self: - d["created"] = d["created"].replace(tzinfo=timezone.utc) + def from_mysql(cls, d: dict) -> Self: + d["created"] = d["created"].replace(tzinfo=UTC) return AuditLog.model_validate(d) diff --git a/generalresearch/models/thl/wallet/cashout_method.py b/generalresearch/models/thl/wallet/cashout_method.py index 59cf721..4757eba 100644 --- a/generalresearch/models/thl/wallet/cashout_method.py +++ b/generalresearch/models/thl/wallet/cashout_method.py @@ -2,7 +2,7 @@ from __future__ import annotations import hashlib import logging -from datetime import datetime, timezone +from datetime import datetime, timezone, UTC from enum import Enum from typing import Any, Literal @@ -16,7 +16,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Self +from typing import Self from generalresearch.currency import USDCent from generalresearch.models.custom_types import ( @@ -145,7 +145,7 @@ class CashoutMethod(CashoutMethodBase): "email associated.", ) last_updated: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) + default_factory=lambda: datetime.now(tz=UTC) ) is_live: bool = Field(default=True) diff --git a/generalresearch/models/thl/wallet/payout.py b/generalresearch/models/thl/wallet/payout.py index 8c78bef..cbb37fe 100644 --- a/generalresearch/models/thl/wallet/payout.py +++ b/generalresearch/models/thl/wallet/payout.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from collections.abc import Collection -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import Any from uuid import uuid4 @@ -50,9 +50,7 @@ class PayoutEvent(BaseModel, validate_assignment=True): # populated from the db and so does not need to be set (there is no # `description` field in event_payout) description: str | None = Field(default=None) - created: AwareDatetimeISO = Field( - default_factory=lambda: datetime.now(tz=timezone.utc) - ) + created: AwareDatetimeISO = Field(default_factory=lambda: datetime.now(tz=UTC)) # In the smallest unit of the currency being transacted. For USD, this # is cents. @@ -159,7 +157,7 @@ class BPPayoutEvent(BaseModel): created: AwareDatetimeISO = Field( description="When the Brokerage Product was paid out", - default_factory=lambda: datetime.now(tz=timezone.utc), + default_factory=lambda: datetime.now(tz=UTC), ) amount: USDCent = Field( diff --git a/generalresearch/pg_helper.py b/generalresearch/pg_helper.py index b9a7d79..b5e124a 100644 --- a/generalresearch/pg_helper.py +++ b/generalresearch/pg_helper.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import timezone +from datetime import UTC, timezone import psycopg from psycopg.adapt import Buffer @@ -24,7 +24,7 @@ class UTCTimestampLoader(TimestampLoader): if dt is None: return None assert dt.tzinfo is None, "expected naive dt" - return dt.replace(tzinfo=timezone.utc) + return dt.replace(tzinfo=UTC) class BPCharLoader(TextLoader): diff --git a/generalresearch/sql_helper.py b/generalresearch/sql_helper.py index 2c526f3..2bbdbe5 100644 --- a/generalresearch/sql_helper.py +++ b/generalresearch/sql_helper.py @@ -130,7 +130,7 @@ def decode_uuids(row: dict[str, Any]) -> dict[str, Any]: class SqlHelper(SqlConnector): - def __init__(self, dsn: Optional[DataBaseDsn] = None, **kwargs): + def __init__(self, dsn: DataBaseDsn | None = None, **kwargs): super().__init__(dsn, **kwargs) def execute_sql_query( @@ -281,7 +281,7 @@ class SqlHelper(SqlConnector): cursor=None, commit=True, primary_key=None, - ) -> Optional[int]: + ) -> int | None: """ Create the item in table `table_name`. In postgresql, `primary_key` needs to be given in order to return the diff --git a/generalresearch/utils/aggregation.py b/generalresearch/utils/aggregation.py index 4023dc9..bd962f1 100644 --- a/generalresearch/utils/aggregation.py +++ b/generalresearch/utils/aggregation.py @@ -2,7 +2,7 @@ from collections import defaultdict from typing import Any, Dict, List -def group_by_year(records: List[Dict], datetime_field: str) -> Dict[int, List[Any]]: +def group_by_year(records: list[dict], datetime_field: str) -> dict[int, list[Any]]: """Memory efficient - processes records one at a time""" by_year = defaultdict(list) diff --git a/generalresearch/utils/copying_cache.py b/generalresearch/utils/copying_cache.py index ea13f69..a1cb37c 100644 --- a/generalresearch/utils/copying_cache.py +++ b/generalresearch/utils/copying_cache.py @@ -1,6 +1,6 @@ +from collections.abc import Callable from copy import deepcopy from functools import wraps -from typing import Callable def deepcopy_return(fn: Callable) -> Callable: diff --git a/generalresearch/utils/enum.py b/generalresearch/utils/enum.py index 14a31de..56706ba 100644 --- a/generalresearch/utils/enum.py +++ b/generalresearch/utils/enum.py @@ -41,7 +41,7 @@ class ReprEnumMeta(EnumMeta): ) -def get_enum_comments(enum_class) -> Dict: +def get_enum_comments(enum_class) -> dict: source = inspect.getsource(enum_class) # Regular expression to match multi-line comments and enum values pattern = re.compile(r"((?:\s*#.*?\n)+)\s*(\w+)\s*=") diff --git a/generalresearch/wall_status_codes/__init__.py b/generalresearch/wall_status_codes/__init__.py index 3a0abb8..1d80924 100644 --- a/generalresearch/wall_status_codes/__init__.py +++ b/generalresearch/wall_status_codes/__init__.py @@ -22,9 +22,9 @@ from generalresearch.wall_status_codes import ( def annotate_status_code( source: Source, ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, Optional[StatusCode1], Optional[str]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1 | None, str | None]: """ :params ext_status_code_1: marketplace-dependent code :params ext_status_code_2: marketplace-dependent code diff --git a/generalresearch/wall_status_codes/cint.py b/generalresearch/wall_status_codes/cint.py index 8042cd2..ecb6219 100644 --- a/generalresearch/wall_status_codes/cint.py +++ b/generalresearch/wall_status_codes/cint.py @@ -6,9 +6,9 @@ from generalresearch.wall_status_codes import lucid def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: return lucid.annotate_status_code( ext_status_code_1=ext_status_code_1, ext_status_code_2=ext_status_code_2, diff --git a/generalresearch/wall_status_codes/dynata.py b/generalresearch/wall_status_codes/dynata.py index 958e06f..69f3e74 100644 --- a/generalresearch/wall_status_codes/dynata.py +++ b/generalresearch/wall_status_codes/dynata.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_name: Dict[str, str] = { +status_codes_name: dict[str, str] = { "0.0": "Unknown", "0.1": "Missing Language", "0.2": "Missing Respondent ID", @@ -51,10 +51,10 @@ status_codes_name: Dict[str, str] = { "5.10": "Daily Limit", } -status_map: Dict[str, Status] = defaultdict( +status_map: dict[str, Status] = defaultdict( lambda: Status.FAIL, **{"1.0": Status.COMPLETE, "1.1": Status.COMPLETE} ) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["1.0", "1.1"], StatusCode1.BUYER_FAIL: ["2.2", "3.2"], StatusCode1.BUYER_QUALITY_FAIL: ["5.1", "5.2"], @@ -88,10 +88,10 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { "5.10", ], } -ext_status_code_map: Dict[str, StatusCode1] = dict() +ext_status_code_map: dict[str, StatusCode1] = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -100,9 +100,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: this is from the callback url params: disposition and status, '.'-joined diff --git a/generalresearch/wall_status_codes/fullcircle.py b/generalresearch/wall_status_codes/fullcircle.py index eda4d9c..f37ffda 100644 --- a/generalresearch/wall_status_codes/fullcircle.py +++ b/generalresearch/wall_status_codes/fullcircle.py @@ -11,7 +11,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_map: Dict[str, str] = { +status_codes_map: dict[str, str] = { "1": "Complete", "2": "Terminate", "3": "Over-quota", @@ -19,7 +19,7 @@ status_codes_map: Dict[str, str] = { } status_map = defaultdict(lambda: Status.FAIL, **{"1": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["1"], StatusCode1.BUYER_FAIL: ["2", "3"], StatusCode1.BUYER_QUALITY_FAIL: ["4"], @@ -29,10 +29,10 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { StatusCode1.PS_FAIL: [], StatusCode1.PS_OVERQUOTA: [], } -ext_status_code_map: Dict[str, StatusCode1] = dict() +ext_status_code_map: dict[str, StatusCode1] = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -41,9 +41,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: this is from the callback url param 's' :params ext_status_code_2: not used diff --git a/generalresearch/wall_status_codes/innovate.py b/generalresearch/wall_status_codes/innovate.py index b650ade..fba3f54 100644 --- a/generalresearch/wall_status_codes/innovate.py +++ b/generalresearch/wall_status_codes/innovate.py @@ -13,7 +13,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_innovate: Dict[str, str] = { +status_codes_innovate: dict[str, str] = { "1": "Complete", "2": "Buyer Fail", "3": "Buyer Over Quota", @@ -29,7 +29,7 @@ status_map = defaultdict( lambda: Status.FAIL, **{"1": Status.COMPLETE, "0": Status.ABANDON, "6": Status.ABANDON}, ) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.BUYER_FAIL: ["2", "3"], StatusCode1.BUYER_QUALITY_FAIL: ["4"], StatusCode1.PS_BLOCKED: [], @@ -43,7 +43,7 @@ for k, v in status_codes_ext_map.items(): for vv in v: ext_status_code_map[status_codes_ext_map.get(vv, vv)] = k -category_innovate: Dict[str, StatusCode1] = { +category_innovate: dict[str, StatusCode1] = { "Selected threat potential score at joblevel not allow the survey": StatusCode1.PS_QUALITY, "OE Validation": StatusCode1.PS_QUALITY, "Unique IP": StatusCode1.PS_DUPLICATE, @@ -78,9 +78,9 @@ category_innovate: Dict[str, StatusCode1] = { def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ Only quality terminate (4 and 8), and PS term (5) return a term_reason (af=). diff --git a/generalresearch/wall_status_codes/lucid.py b/generalresearch/wall_status_codes/lucid.py index c0098e6..3ce89b0 100644 --- a/generalresearch/wall_status_codes/lucid.py +++ b/generalresearch/wall_status_codes/lucid.py @@ -9,7 +9,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -mp_codes: Dict[str, str] = { +mp_codes: dict[str, str] = { "-6": "Pre-Client Intermediary Page Drop Off", "-5": "Failure in the Post Answer Behavior", "-1": "Failure to Load the Lucid Marketplace", @@ -54,7 +54,7 @@ mp_codes: Dict[str, str] = { } # todo: finish, there's a bunch more -client_status_map: Dict[str, StatusCode1] = { +client_status_map: dict[str, StatusCode1] = { "30": StatusCode1.BUYER_QUALITY_FAIL, "33": StatusCode1.BUYER_QUALITY_FAIL, "34": StatusCode1.BUYER_QUALITY_FAIL, @@ -62,7 +62,7 @@ client_status_map: Dict[str, StatusCode1] = { } status_map = defaultdict(lambda: Status.FAIL, **{"s": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: [], StatusCode1.BUYER_FAIL: ["3"], StatusCode1.BUYER_QUALITY_FAIL: [], @@ -102,10 +102,10 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { StatusCode1.PS_OVERQUOTA: ["40", "41", "42"], } -ext_status_code_map: Dict[str, StatusCode1] = dict() +ext_status_code_map: dict[str, StatusCode1] = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -115,9 +115,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: this indicates which callback url was hit. possible values {'s', *anything else*} :params ext_status_code_2: this is from the callback url params: InitialStatus diff --git a/generalresearch/wall_status_codes/morning.py b/generalresearch/wall_status_codes/morning.py index 318d1b2..42e34b3 100644 --- a/generalresearch/wall_status_codes/morning.py +++ b/generalresearch/wall_status_codes/morning.py @@ -17,7 +17,7 @@ timeout: The respondent completed the survey after the timeout period had expire in_progress: The respondent interview session is still in progress, such as in the prescreener or survey. """ -short_code_to_status_codes_morning: Dict[str, str] = { +short_code_to_status_codes_morning: dict[str, str] = { "att_che": "attention_check", "banned": "banned", "bid_clo": "bid_closed", @@ -54,7 +54,7 @@ short_code_to_status_codes_morning: Dict[str, str] = { } status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["complete"], StatusCode1.BUYER_FAIL: [ "in_survey_failure", @@ -97,10 +97,10 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { "quota_invalid_for_bid", ], } -ext_status_code_map: Dict[str, StatusCode1] = dict() +ext_status_code_map: dict[str, StatusCode1] = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -109,9 +109,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: from callback url params: &sti={{status_id}} :params ext_status_code_2: from callback url params: &sdi={{status_detail_id}} diff --git a/generalresearch/wall_status_codes/pollfish.py b/generalresearch/wall_status_codes/pollfish.py index 361785d..1ae732d 100644 --- a/generalresearch/wall_status_codes/pollfish.py +++ b/generalresearch/wall_status_codes/pollfish.py @@ -1,9 +1,9 @@ from collections import defaultdict -from typing import Any, Dict, List, Optional, Tuple +from typing import Any from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_map: Dict[str, str] = { +status_codes_map: dict[str, str] = { "quo_ful": "quota_full", "sur_clo": "survey_closed", "profilin": "profiling", @@ -29,7 +29,7 @@ status_codes_map: Dict[str, str] = { "complete": "complete", } status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["complete"], StatusCode1.BUYER_FAIL: ["third_party_termination", "screenout"], StatusCode1.BUYER_QUALITY_FAIL: [ @@ -61,7 +61,7 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { ext_status_code_map = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -70,9 +70,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: from callback url params: &sti={{status_id}} :params ext_status_code_2: from callback url params: &sdi={{status_detail_id}} diff --git a/generalresearch/wall_status_codes/precision.py b/generalresearch/wall_status_codes/precision.py index ffeaeca..d3c4471 100644 --- a/generalresearch/wall_status_codes/precision.py +++ b/generalresearch/wall_status_codes/precision.py @@ -11,7 +11,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_precision: Dict[str, str] = { +status_codes_precision: dict[str, str] = { "10": "Complete", "20": "Client Terminate", "21": "PS Terminate", @@ -46,7 +46,7 @@ status_codes_precision: Dict[str, str] = { "80": "Final Complete", } status_map = defaultdict(lambda: Status.FAIL, **{"s": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["10"], StatusCode1.BUYER_FAIL: ["20", "30"], StatusCode1.BUYER_QUALITY_FAIL: ["60"], @@ -76,7 +76,7 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { ext_status_code_map = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -85,9 +85,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: from callback url params: status :params ext_status_code_2: from callback url params: code diff --git a/generalresearch/wall_status_codes/prodege.py b/generalresearch/wall_status_codes/prodege.py index a1e25ea..aaac376 100644 --- a/generalresearch/wall_status_codes/prodege.py +++ b/generalresearch/wall_status_codes/prodege.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 status_map = defaultdict(lambda: Status.FAIL, **{"1": Status.COMPLETE}) -status_code_map: Dict[StatusCode1, List[str]] = { +status_code_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: [], StatusCode1.BUYER_FAIL: ["1", "2"], StatusCode1.BUYER_QUALITY_FAIL: ["10", "12"], @@ -34,7 +34,7 @@ status_code_map: Dict[StatusCode1, List[str]] = { status_class = dict() for k, v in status_code_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -43,9 +43,9 @@ for k, v in status_code_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: status from redirect url :params ext_status_code_2: termreason from redirect url diff --git a/generalresearch/wall_status_codes/repdata.py b/generalresearch/wall_status_codes/repdata.py index c502532..8b690b5 100644 --- a/generalresearch/wall_status_codes/repdata.py +++ b/generalresearch/wall_status_codes/repdata.py @@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_name: Dict[str, str] = { +status_codes_name: dict[str, str] = { "2": "Search Failed", "3": "Activity Failed", "4": "Review Failed", @@ -26,7 +26,7 @@ status_codes_name: Dict[str, str] = { "6003": "In-Survey maximum exceeded (Research Desk)", } # See: 02, and 13 are de-dupes -rd_threat_name: Dict[str, str] = { +rd_threat_name: dict[str, str] = { "02": "Duplicate entrant into survey", "03": "Emulator Usage", "04": "VPN usage detected", @@ -47,7 +47,7 @@ rd_threat_name: Dict[str, str] = { } status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) -status_code_map: Dict[StatusCode1, List[str]] = { +status_code_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["1000"], StatusCode1.BUYER_FAIL: ["2000", "4000"], StatusCode1.BUYER_QUALITY_FAIL: ["3000"], @@ -61,7 +61,7 @@ status_code_map: Dict[StatusCode1, List[str]] = { status_class = dict() for k, v in status_code_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -70,9 +70,9 @@ for k, v in status_code_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: the redirect urls category (as defined in url param 549f3710b) {'term', 'overquota', 'fraud', 'complete'} diff --git a/generalresearch/wall_status_codes/sago.py b/generalresearch/wall_status_codes/sago.py index 9c8710f..028192b 100644 --- a/generalresearch/wall_status_codes/sago.py +++ b/generalresearch/wall_status_codes/sago.py @@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_schlesinger: Dict[str, str] = { +status_codes_schlesinger: dict[str, str] = { "1": "Complete", "2": "Buyer Fail", "3": "Buyer Fail", @@ -20,7 +20,7 @@ status_codes_schlesinger: Dict[str, str] = { "11": "Abandon", # really it is "Buyer Abandon" } -status_reason_name: Dict[str, str] = { +status_reason_name: dict[str, str] = { "1": "Not a Unique Sample Cube User", "4": "GeoIP - wrong country", "7": "Duplicate - not a unique IP", @@ -121,7 +121,7 @@ status_map = defaultdict( lambda: Status.FAIL, **{"1": Status.COMPLETE, "0": Status.ABANDON} ) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["48"], StatusCode1.BUYER_FAIL: ["16", "29", "49", "50", "78", "114", "110", "114"], StatusCode1.BUYER_QUALITY_FAIL: ["26", "52", "68", "81", "84"], @@ -167,10 +167,10 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { StatusCode1.PS_FAIL: ["7", "29", "36", "47", "56", "58", "64"], StatusCode1.PS_OVERQUOTA: ["29", "46", "33", "31"], } -ext_status_code_map: Dict[str, StatusCode1] = dict() +ext_status_code_map: dict[str, StatusCode1] = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -179,9 +179,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: from callback url params: scstatus :params ext_status_code_2: from callback url params: scsecuritystatus diff --git a/generalresearch/wall_status_codes/spectrum.py b/generalresearch/wall_status_codes/spectrum.py index 9e9e0a2..cf200a9 100644 --- a/generalresearch/wall_status_codes/spectrum.py +++ b/generalresearch/wall_status_codes/spectrum.py @@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional, Tuple from generalresearch.models.thl.definitions import Status, StatusCode1 -status_codes_spectrum: Dict[str, str] = { +status_codes_spectrum: dict[str, str] = { "11": "PS Drop", "12": "PS Quota Full Core", "13": "PS Termination Core", @@ -80,7 +80,7 @@ status_codes_spectrum: Dict[str, str] = { "88": "PS_Supplier_Allocation_Throttle", } status_map = defaultdict(lambda: Status.FAIL, **{"21": Status.COMPLETE}) -status_codes_ext_map: Dict[StatusCode1, List[str]] = { +status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["21"], StatusCode1.BUYER_FAIL: ["16", "17", "18", "19", "30", "59", "84"], StatusCode1.BUYER_QUALITY_FAIL: ["20", "31"], @@ -143,7 +143,7 @@ status_codes_ext_map: Dict[StatusCode1, List[str]] = { ext_status_code_map = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[str] + v: list[str] for vv in v: vv: str @@ -152,9 +152,9 @@ for k, v in status_codes_ext_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[Any]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, Any | None]: """ :params ext_status_code_1: from url params: ps_rstatus https://purespectrum.atlassian.net/wiki/spaces/PA/pages/33613201/Minimizing+Clickwaste+with+ps+rstatus diff --git a/generalresearch/wall_status_codes/wxet.py b/generalresearch/wall_status_codes/wxet.py index e7cf67d..1b7c514 100644 --- a/generalresearch/wall_status_codes/wxet.py +++ b/generalresearch/wall_status_codes/wxet.py @@ -8,10 +8,10 @@ from generalresearch.wxet.models.definitions import ( WXETStatusCode2, ) -status_map: Dict[WXETStatus, Status] = defaultdict( +status_map: dict[WXETStatus, Status] = defaultdict( lambda: Status.FAIL, **{WXETStatus.COMPLETE: Status.COMPLETE} ) -status_codes_ext_map: Dict[StatusCode1, List[WXETStatusCode1]] = { +status_codes_ext_map: dict[StatusCode1, list[WXETStatusCode1]] = { StatusCode1.COMPLETE: [WXETStatusCode1.COMPLETE], StatusCode1.BUYER_FAIL: [ WXETStatusCode1.BUYER_DUPLICATE, @@ -33,13 +33,13 @@ status_codes_ext_map: Dict[StatusCode1, List[WXETStatusCode1]] = { ext_status_code_map = dict() for k, v in status_codes_ext_map.items(): k: StatusCode1 - v: List[WXETStatusCode1] + v: list[WXETStatusCode1] for vv in v: vv: WXETStatusCode1 ext_status_code_map[vv] = k -status_code2_map: Dict[StatusCode1, List[WXETStatusCode2]] = { +status_code2_map: dict[StatusCode1, list[WXETStatusCode2]] = { StatusCode1.PS_QUALITY: [], StatusCode1.PS_DUPLICATE: [ WXETStatusCode2.WORKER_INELIGIBLE, @@ -67,9 +67,9 @@ for k, v in status_code2_map.items(): def annotate_status_code( ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, -) -> Tuple[Status, StatusCode1, Optional[WXETStatusCode2]]: + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, +) -> tuple[Status, StatusCode1, WXETStatusCode2 | None]: """ :params ext_status_code_1: WXETStatus :params ext_status_code_2: WXETStatusCode1 diff --git a/generalresearch/wxet/models/definitions.py b/generalresearch/wxet/models/definitions.py index 8d0d6b8..b08e178 100644 --- a/generalresearch/wxet/models/definitions.py +++ b/generalresearch/wxet/models/definitions.py @@ -166,8 +166,8 @@ class WXETStatusCode2(int, Enum, metaclass=ReprEnumMeta): def check_wxet_status_consistent( status: WXETStatus, - status_code_1: Optional[WXETStatusCode1] = None, - status_code_2: Optional[WXETStatusCode2] = None, + status_code_1: WXETStatusCode1 | None = None, + status_code_2: WXETStatusCode2 | None = None, ) -> bool: """ Raises an AssertionError if inconsistent @@ -203,13 +203,13 @@ def check_wxet_status_consistent( def check_wxet_adjusted_status_attempt_consistent( status: WXETStatus, - status_code_1: Optional[WXETStatusCode1] = None, - cpi: Optional[USDMill] = None, - adjusted_status: Optional[WXETAdjustedStatus] = None, - adjusted_cpi: Optional[USDMill] = None, - new_adjusted_status: Optional[WXETAdjustedStatus] = None, - new_adjusted_cpi: Optional[USDMill] = None, -) -> Tuple[bool, str]: + status_code_1: WXETStatusCode1 | None = None, + cpi: USDMill | None = None, + adjusted_status: WXETAdjustedStatus | None = None, + adjusted_cpi: USDMill | None = None, + new_adjusted_status: WXETAdjustedStatus | None = None, + new_adjusted_cpi: USDMill | None = None, +) -> tuple[bool, str]: """ Raises an AssertionError if inconsistent. - status, status_code_1, adjusted_status, adjusted_cpi, cpi are the attempt's CURRENT values @@ -233,12 +233,12 @@ def check_wxet_adjusted_status_attempt_consistent( def _check_wxet_adjusted_status_attempt_consistent( status: WXETStatus, - status_code_1: Optional[WXETStatusCode1] = None, - cpi: Optional[USDMill] = None, - adjusted_status: Optional[WXETAdjustedStatus] = None, - adjusted_cpi: Optional[USDMill] = None, - new_adjusted_status: Optional[WXETAdjustedStatus] = None, - new_adjusted_cpi: Optional[USDMill] = None, + status_code_1: WXETStatusCode1 | None = None, + cpi: USDMill | None = None, + adjusted_status: WXETAdjustedStatus | None = None, + adjusted_cpi: USDMill | None = None, + new_adjusted_status: WXETAdjustedStatus | None = None, + new_adjusted_cpi: USDMill | None = None, ) -> None: """ Raises an AssertionError if inconsistent. @@ -297,8 +297,8 @@ def _check_wxet_adjusted_status_attempt_consistent( def _check_wxet_adjusted_status_consistent( - adjusted_status: Optional[WXETAdjustedStatus] = None, - adjusted_cpi: Optional[USDMill] = None, + adjusted_status: WXETAdjustedStatus | None = None, + adjusted_cpi: USDMill | None = None, ) -> None: """ Raises an AssertionError if inconsistent. diff --git a/generalresearch/wxet/models/finish_type.py b/generalresearch/wxet/models/finish_type.py index af60fe6..a57dce8 100644 --- a/generalresearch/wxet/models/finish_type.py +++ b/generalresearch/wxet/models/finish_type.py @@ -33,7 +33,7 @@ class FinishType(str, Enum, metaclass=ReprEnumMeta): FAIL = "fail" @property - def finish_statuses(self) -> Set[Optional[WXETStatus]]: + def finish_statuses(self) -> set[WXETStatus | None]: """For this particular FinishType, what are the different WXETStatus values that are consider """ @@ -64,9 +64,9 @@ class FinishType(str, Enum, metaclass=ReprEnumMeta): def is_a_finish( - status: Optional[WXETStatus], - status_code_1: Optional[WXETStatusCode1], - finish_type: Optional[FinishType], + status: WXETStatus | None, + status_code_1: WXETStatusCode1 | None, + finish_type: FinishType | None, ) -> bool: """Determines if a wall event should be considered a finish or not. diff --git a/test_utils/conftest.py b/test_utils/conftest.py index 378b9cc..ffe458c 100644 --- a/test_utils/conftest.py +++ b/test_utils/conftest.py @@ -6,16 +6,17 @@ import stat import subprocess import sys import tempfile -from datetime import datetime, timedelta, timezone +from collections.abc import Callable, Generator +from datetime import UTC, datetime, timedelta, timezone from os.path import join as pjoin from pathlib import Path -from typing import Callable, Generator from uuid import uuid4 import pytest from _pytest.config import Config from dotenv import load_dotenv from pydantic import MariaDBDsn, PostgresDsn, TypeAdapter +from pytest import TempPathFactory from generalresearch.config import GRLBaseSettings from generalresearch.currency import USDCent @@ -93,7 +94,7 @@ def postgres_instance(settings: GRLBaseSettings) -> Generator[PostgresDsn]: from psycopg import connect from psycopg.sql import SQL, Identifier - now = datetime.now(timezone.utc) + now = datetime.now(UTC) ts: str = now.strftime("%Y-%m-%d") db_name = f"unittest-{ts}-{uuid4().hex[:6]}" @@ -152,38 +153,48 @@ def postgres_instance_host( yield value -# @pytest.fixture(scope="session") -# def git_key_path(settings: GRLBaseSettings) -> Path: -# return Path('/tmp/') - - @pytest.fixture(scope="session") def git_key_path( + tmp_path_factory: TempPathFactory, settings: GRLBaseSettings, ) -> Generator[Path]: + # We are using the tmp_path_factory because unlike the tmp_path (which + # is function scoped), this is session scoped. - assert settings.git_creds - with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix="_id_rsa") as f: - f.write(settings.git_creds) - key_path = f.name - - os.chmod(key_path, stat.S_IRUSR | stat.S_IWUSR) + assert settings.git_creds, "Must define key to download alternative models" + fn = tmp_path_factory.mktemp("keys") / "git_creds" + fn.write_text(settings.git_creds, encoding="utf-8") + os.chmod(fn, stat.S_IRUSR | stat.S_IWUSR) - yield Path(key_path) + yield Path(fn) - os.unlink(key_path) + os.unlink(fn) @pytest.fixture(scope="session") -def gr_repo(git_key_path: Path) -> Callable[..., Path]: +def gr_repo( + git_key_path: Path, + tmp_path_factory: TempPathFactory, +) -> Callable[..., Path | None]: repo_url = "ssh://code.g-r-l.com/general-research/gr-carer.git" - repo_path = Path("/tmp/gr-carer") + + _ran = {} + if _ran.get(repo_url, False): + print(f"Already ran django_db_factory.{repo_url}") + return + + _ran[repo_url] = True + + fn = tmp_path_factory.mktemp("repos") + repo_path = fn / "gr-carer" + repo_path.mkdir(parents=True, exist_ok=True) def _inner() -> Path: + ssh_cmd = ( f"ssh -i {git_key_path} " "-o IdentitiesOnly=yes " - "-o StrictHostKeyChecking=no " # or accept-new, see note below + "-o StrictHostKeyChecking=no " ) env = {"GIT_SSH_COMMAND": ssh_cmd} @@ -196,6 +207,11 @@ def gr_repo(git_key_path: Path) -> Callable[..., Path]: env=env, ) + result = subprocess.run( + ["cat", git_key_path], capture_output=True, text=True, check=False + ) + print(repr(result.stdout)) + return repo_path return _inner @@ -206,21 +222,29 @@ def django_db_factory( postgres_instance: PostgresDsn, postgres_instance_dict: PostgresDict, gr_repo: Callable[..., Path], -) -> Callable[..., PostgresDsn]: +) -> Callable[..., PostgresDsn | None]: + + _ran = {} import django + from django.apps import apps from django.conf import settings as django_settings from django.core.management import call_command - def _inner(django_project: str = "generalresearch.thl_django"): + def _inner( + django_project: str = "generalresearch.thl_django", + ) -> PostgresDsn | None: + + if _ran.get(django_project, False): + print(f"Already ran django_db_factory.{django_project}") + return + _ran[django_project] = True if "gr" in django_project: # We need model files that are NOT in this repo. gr_path = gr_repo() sys.path.insert(0, str(gr_path)) - print(sys.path) - # 1. Bootstrapping Django settings if not django_settings.configured: django_settings.configure( @@ -242,10 +266,11 @@ def django_db_factory( ) django.setup() - # for model in apps.get_models(): - # print(f"Discovered model: {model._meta.label}") + for model in apps.get_models(): + print(f"Discovered model: {model._meta.label}") # 2. Run migrations directly during fixture activation + call_command("makemigrations", "gr", interactive=False) call_command("migrate") # 3. Return the Dsn so the factory gives a way to connect @@ -276,37 +301,37 @@ def spectrum_rw(settings: GRLBaseSettings) -> SqlHelper: @pytest.fixture def start() -> datetime: - return datetime(year=1900, month=1, day=1, tzinfo=timezone.utc) + return datetime(year=1900, month=1, day=1, tzinfo=UTC) @pytest.fixture def utc_now() -> datetime: - return datetime.now(tz=timezone.utc) + return datetime.now(tz=UTC) @pytest.fixture def utc_hour_ago() -> datetime: - return datetime.now(tz=timezone.utc) - timedelta(hours=1) + return datetime.now(tz=UTC) - timedelta(hours=1) @pytest.fixture def utc_day_ago() -> datetime: - return datetime.now(tz=timezone.utc) - timedelta(hours=24) + return datetime.now(tz=UTC) - timedelta(hours=24) @pytest.fixture def utc_90days_ago() -> datetime: - return datetime.now(tz=timezone.utc) - timedelta(days=90) + return datetime.now(tz=UTC) - timedelta(days=90) @pytest.fixture def utc_60days_ago() -> datetime: - return datetime.now(tz=timezone.utc) - timedelta(days=60) + return datetime.now(tz=UTC) - timedelta(days=60) @pytest.fixture def utc_30days_ago() -> datetime: - return datetime.now(tz=timezone.utc) - timedelta(days=30) + return datetime.now(tz=UTC) - timedelta(days=30) # === Clean up === @@ -322,7 +347,7 @@ def delete_df_collection( DFCollectionType, ) - def _inner(coll: "DFCollection"): + def _inner(coll: DFCollection): match coll.data_type: case DFCollectionType.LEDGER: for table in [ diff --git a/test_utils/grliq/conftest.py b/test_utils/grliq/conftest.py index e8175a5..7665b52 100644 --- a/test_utils/grliq/conftest.py +++ b/test_utils/grliq/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone -from typing import Callable +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from uuid import uuid4 import pytest @@ -83,7 +83,7 @@ def grliq_data() -> GrlIqData: g.id = None g.uuid = uuid4().hex - g.created_at = datetime.now(tz=timezone.utc) + g.created_at = datetime.now(tz=UTC) g.timestamp = g.created_at - timedelta(seconds=10) return g @@ -117,7 +117,7 @@ def grliq_data_factory(grliq_dm: GrlIqDataManager) -> Callable[..., GrlIqData]: product_user_id = product_user_id or uuid4().hex uuid = uuid or uuid4().hex mid = mid or uuid4().hex - created_at = created_at or datetime.now(tz=timezone.utc) + created_at = created_at or datetime.now(tz=UTC) res["data"].product_id = product_id res["data"].product_user_id = product_user_id diff --git a/test_utils/incite/collections/conftest.py b/test_utils/incite/collections/conftest.py index 88eef72..631bb7b 100644 --- a/test_utils/incite/collections/conftest.py +++ b/test_utils/incite/collections/conftest.py @@ -1,7 +1,8 @@ from __future__ import annotations +from collections.abc import Callable from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING import pytest diff --git a/test_utils/incite/conftest.py b/test_utils/incite/conftest.py index 12e57c5..87ea7ae 100644 --- a/test_utils/incite/conftest.py +++ b/test_utils/incite/conftest.py @@ -1,11 +1,12 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from os.path import join as pjoin from pathlib import Path from random import choice as randchoice from shutil import rmtree -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -166,7 +167,7 @@ def incite_item_factory( for _ in range(5): item_time = fake.date_time_between( - start_date=item.start, end_date=item.finish, tzinfo=timezone.utc + start_date=item.start, end_date=item.finish, tzinfo=UTC ) match data_type: diff --git a/test_utils/incite/mergers/conftest.py b/test_utils/incite/mergers/conftest.py index e9970c2..c0f0bcf 100644 --- a/test_utils/incite/mergers/conftest.py +++ b/test_utils/incite/mergers/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations +from collections.abc import Callable from datetime import datetime, timedelta -from typing import Callable import pytest diff --git a/test_utils/managers/conftest.py b/test_utils/managers/conftest.py index d2e5d20..4dacb29 100644 --- a/test_utils/managers/conftest.py +++ b/test_utils/managers/conftest.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable +from collections.abc import Callable import pytest diff --git a/test_utils/managers/gr/conftest.py b/test_utils/managers/gr/conftest.py index 37da164..4da8fe3 100644 --- a/test_utils/managers/gr/conftest.py +++ b/test_utils/managers/gr/conftest.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable +from collections.abc import Callable import pytest import redis.asyncio as redis_async diff --git a/test_utils/managers/thl/conftest.py b/test_utils/managers/thl/conftest.py index 5b70961..21b2007 100644 --- a/test_utils/managers/thl/conftest.py +++ b/test_utils/managers/thl/conftest.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable +from collections.abc import Callable import pytest from pydantic import PostgresDsn diff --git a/test_utils/managers/upk/conftest.py b/test_utils/managers/upk/conftest.py index d8f956c..7eabee1 100644 --- a/test_utils/managers/upk/conftest.py +++ b/test_utils/managers/upk/conftest.py @@ -1,4 +1,4 @@ -from typing import Callable, Generator +from collections.abc import Callable, Generator import pytest diff --git a/test_utils/models/conftest.py b/test_utils/models/conftest.py index 93e2f44..3a9e45c 100644 --- a/test_utils/models/conftest.py +++ b/test_utils/models/conftest.py @@ -1,10 +1,11 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from random import choice as randchoice from random import randint -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -124,7 +125,7 @@ def wall_factory(wall_manager: WallManager) -> Callable[..., Wall]: ) -> Wall: assert session.started <= datetime.now( - tz=timezone.utc + tz=UTC ), "Session can't start in the future" if session.wall_events: diff --git a/test_utils/models/contest/conftest.py b/test_utils/models/contest/conftest.py index bfbe9f8..0060946 100644 --- a/test_utils/models/contest/conftest.py +++ b/test_utils/models/contest/conftest.py @@ -1,14 +1,13 @@ from __future__ import annotations -from datetime import datetime, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timezone from decimal import Decimal -from typing import Callable from uuid import uuid4 import pytest from pytest import FixtureRequest as Request - from generalresearch.currency import USDCent from generalresearch.managers.thl.contest_manager import ContestManager from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager @@ -136,7 +135,7 @@ def milestone_contest_create() -> MilestoneContestCreate: ), ], end_condition=MilestoneContestEndCondition( - ends_at=datetime(year=2030, month=1, day=1, tzinfo=timezone.utc), + ends_at=datetime(year=2030, month=1, day=1, tzinfo=UTC), max_winners=5, ), entry_trigger=ContestEntryTrigger.TASK_COMPLETE, diff --git a/test_utils/models/gr/conftest.py b/test_utils/models/gr/conftest.py index df97306..90b86aa 100644 --- a/test_utils/models/gr/conftest.py +++ b/test_utils/models/gr/conftest.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable +from collections.abc import Callable from uuid import uuid4 import pytest diff --git a/test_utils/models/ledger/conftest.py b/test_utils/models/ledger/conftest.py index 14a7465..b428468 100644 --- a/test_utils/models/ledger/conftest.py +++ b/test_utils/models/ledger/conftest.py @@ -1,9 +1,10 @@ from __future__ import annotations +from collections.abc import Callable from datetime import datetime from decimal import Decimal from random import randint -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py index bebc691..cabd8dc 100644 --- a/test_utils/models/network/conftest.py +++ b/test_utils/models/network/conftest.py @@ -1,5 +1,5 @@ import os -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from uuid import uuid4 import pytest @@ -90,7 +90,7 @@ def rdns_result(dig_raw_output: str) -> RDNSResult: def rdns_run(rdns_result: RDNSResult, scan_group_id: str): r = rdns_result ip = "45.33.32.156" - utc_now = datetime.now(tz=timezone.utc) + utc_now = datetime.now(tz=UTC) config = RDNSRunCommand(command="dig", options=RDNSRunCommandOptions(ip=ip)) return RDNSRun( tool_version="1.2.3", @@ -121,7 +121,7 @@ def mtr_result(mtr_raw_output: str) -> MTRResult: @pytest.fixture(scope="session") def mtr_run(mtr_result: MTRResult, scan_group_id: str): r = mtr_result - utc_now = datetime.now(tz=timezone.utc) + utc_now = datetime.now(tz=UTC) config = MTRRunCommand( command="mtr", options=MTRRunCommandOptions( diff --git a/test_utils/models/thl/conftest.py b/test_utils/models/thl/conftest.py index cf8d2fa..a2adcce 100644 --- a/test_utils/models/thl/conftest.py +++ b/test_utils/models/thl/conftest.py @@ -1,11 +1,12 @@ from __future__ import annotations -from datetime import datetime, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timezone from decimal import ROUND_DOWN, Decimal from random import choice as rand_choice from random import choice as rchoice from random import randint, random -from typing import Any, Callable +from typing import Any from uuid import uuid4 import faker @@ -105,9 +106,9 @@ def wall_factory( user_id = user_id or fake.random_int(min=1, max=2_147_483_648) started = started or fake.date_time_between( - start_date=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), - end_date=datetime.now(tz=timezone.utc), - tzinfo=timezone.utc, + start_date=datetime(year=1900, month=1, day=1, tzinfo=UTC), + end_date=datetime.now(tz=UTC), + tzinfo=UTC, ) if session_id is None: @@ -199,9 +200,9 @@ def session_factory(session_manager: SessionManager): ) -> Session: """To be used in tests, where we don't care about certain fields""" started = started or fake.date_time_between( - start_date=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), - end_date=datetime(year=2000, month=1, day=1, tzinfo=timezone.utc), - tzinfo=timezone.utc, + start_date=datetime(year=1900, month=1, day=1, tzinfo=UTC), + end_date=datetime(year=2000, month=1, day=1, tzinfo=UTC), + tzinfo=UTC, ) user = user or User( user_id=fake.random_int(min=1, max=2_147_483_648), uuid=uuid4().hex diff --git a/test_utils/spectrum/conftest.py b/test_utils/spectrum/conftest.py index 0afc3f5..9c067d3 100644 --- a/test_utils/spectrum/conftest.py +++ b/test_utils/spectrum/conftest.py @@ -1,6 +1,6 @@ import logging import time -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import TYPE_CHECKING import pytest @@ -19,7 +19,7 @@ if TYPE_CHECKING: @pytest.fixture(scope="session") -def spectrum_rw(settings: "GRLBaseSettings") -> SqlHelper: +def spectrum_rw(settings: GRLBaseSettings) -> SqlHelper: logging.info(f"{settings.spectrum_rw_db=}") assert settings.spectrum_rw_db is not None @@ -49,11 +49,11 @@ def spectrum_survey_manager(spectrum_rw: SqlHelper) -> SpectrumSurveyManager: def setup_spectrum_surveys( spectrum_rw: SqlHelper, spectrum_survey_manager, spectrum_criteria_manager ) -> None: - now = datetime.now(timezone.utc) + now = datetime.now(UTC) # make sure these example surveys exist in db surveys = [SpectrumSurvey.model_validate_json(x) for x in SURVEYS_JSON] for s in surveys: - s.modified_api = datetime.now(tz=timezone.utc) + s.modified_api = datetime.now(tz=UTC) spectrum_survey_manager.create_or_update(surveys) spectrum_criteria_manager.update(CONDITIONS) @@ -66,10 +66,10 @@ def setup_spectrum_surveys( ["687", "GRL", "x", "x", "x", "x"], commit=True, ) - supplier687_pk = spectrum_rw.execute_sql_query( - f""" - select id from `{spectrum_rw.db}`.spectrum_supplier where supplier_id = '687'""" - )[0]["id"] + supplier687_pk = spectrum_rw.execute_sql_query(f""" + select id from `{spectrum_rw.db}`.spectrum_supplier where supplier_id = '687'""")[ + 0 + ]["id"] conn = spectrum_rw.make_connection() c = conn.cursor() c.executemany( diff --git a/tests/grliq/models/test_forensic_data.py b/tests/grliq/models/test_forensic_data.py index 4fbf962..a901dc3 100644 --- a/tests/grliq/models/test_forensic_data.py +++ b/tests/grliq/models/test_forensic_data.py @@ -9,16 +9,16 @@ if TYPE_CHECKING: class TestGrlIqData: - def test_supported_fonts(self, grliq_data: "GrlIqData"): + def test_supported_fonts(self, grliq_data: GrlIqData): s = grliq_data.supported_fonts_binary assert len(s) == 1043 assert "Ubuntu" in grliq_data.supported_fonts - def test_battery(self, grliq_data: "GrlIqData"): + def test_battery(self, grliq_data: GrlIqData): assert not grliq_data.battery_charging assert grliq_data.battery_level == 0.41 - def test_base(self, grliq_data: "GrlIqData"): + def test_base(self, grliq_data: GrlIqData): from generalresearch.grliq.models.forensic_data import Platform assert grliq_data.timezone == "America/Los_Angeles" @@ -41,7 +41,7 @@ class TestGrlIqData: # Testing things that will cause a validation error, should only be # because something is "corrupt", not b/c the user is a baddie - def test_corrupt(self, grliq_data: "GrlIqData"): + def test_corrupt(self, grliq_data: GrlIqData): """Test for timestamp and timezone offset mismatch validation.""" from generalresearch.grliq.models.forensic_data import GrlIqData diff --git a/tests/incite/collections/test_df_collection_base.py b/tests/incite/collections/test_df_collection_base.py index 31d1720..c3c64e4 100644 --- a/tests/incite/collections/test_df_collection_base.py +++ b/tests/incite/collections/test_df_collection_base.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import TYPE_CHECKING import pandas as pd @@ -24,7 +24,7 @@ class TestDFCollectionBase: """ - def test_init(self, mnt_filepath: "GRLDatasets", df_coll_type: DFCollectionType): + def test_init(self, mnt_filepath: GRLDatasets, df_coll_type: DFCollectionType): """Try to initialize the DFCollection with various invalid parameters""" with pytest.raises(expected_exception=ValueError) as cm: DFCollection(archive_path=mnt_filepath.data_src) @@ -46,11 +46,11 @@ class TestDFCollectionBase: class TestDFCollectionBaseProperties: @pytest.mark.skip - def test_df_collection_items(self, mnt_filepath: "GRLDatasets", df_coll_type): + def test_df_collection_items(self, mnt_filepath: GRLDatasets, df_coll_type): instance = DFCollection( data_type=df_coll_type, - start=datetime(year=1800, month=1, day=1, tzinfo=timezone.utc), - finished=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=1800, month=1, day=1, tzinfo=UTC), + finished=datetime(year=1900, month=1, day=1, tzinfo=UTC), offset="100d", archive_path=mnt_filepath.archive_path(enum_type=df_coll_type), ) @@ -58,11 +58,11 @@ class TestDFCollectionBaseProperties: assert len(instance.interval_range) == len(instance.items) assert len(instance.items) == 366 - def test_df_collection_progress(self, mnt_filepath: "GRLDatasets", df_coll_type): + def test_df_collection_progress(self, mnt_filepath: GRLDatasets, df_coll_type): instance = DFCollection( data_type=df_coll_type, - start=datetime(year=1800, month=1, day=1, tzinfo=timezone.utc), - finished=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=1800, month=1, day=1, tzinfo=UTC), + finished=datetime(year=1900, month=1, day=1, tzinfo=UTC), offset="100d", archive_path=mnt_filepath.archive_path(enum_type=df_coll_type), ) @@ -71,7 +71,7 @@ class TestDFCollectionBaseProperties: assert isinstance(instance.progress, pd.DataFrame) assert instance.progress.shape == (366, 6) - def test_df_collection_schema(self, mnt_filepath: "GRLDatasets", df_coll_type): + def test_df_collection_schema(self, mnt_filepath: GRLDatasets, df_coll_type): instance1 = DFCollection( data_type=DFCollectionType.WALL, archive_path=mnt_filepath.data_src ) @@ -88,12 +88,12 @@ class TestDFCollectionBaseProperties: class TestDFCollectionBaseMethods: @pytest.mark.skip - def test_initial_load(self, mnt_filepath: "GRLDatasets", thl_web_rr): + def test_initial_load(self, mnt_filepath: GRLDatasets, thl_web_rr): instance = DFCollection( pg_config=thl_web_rr, data_type=DFCollectionType.USER, - start=datetime(year=2022, month=1, day=1, minute=0, tzinfo=timezone.utc), - finished=datetime(year=2022, month=1, day=1, minute=5, tzinfo=timezone.utc), + start=datetime(year=2022, month=1, day=1, minute=0, tzinfo=UTC), + finished=datetime(year=2022, month=1, day=1, minute=5, tzinfo=UTC), offset="2min", archive_path=mnt_filepath.data_src, ) diff --git a/tests/incite/collections/test_df_collection_item_base.py b/tests/incite/collections/test_df_collection_item_base.py index 136d234..8cf719d 100644 --- a/tests/incite/collections/test_df_collection_item_base.py +++ b/tests/incite/collections/test_df_collection_item_base.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from typing import TYPE_CHECKING import pytest @@ -19,12 +19,12 @@ df_collection_types = [e for e in DFCollectionType if e is not DFCollectionType. @pytest.mark.parametrize("df_coll_type", df_collection_types) class TestDFCollectionItemBase: - def test_init(self, mnt_filepath: "GRLDatasets", df_coll_type): + def test_init(self, mnt_filepath: GRLDatasets, df_coll_type): collection = DFCollection( data_type=df_coll_type, offset="100d", - start=datetime(year=1800, month=1, day=1, tzinfo=timezone.utc), - finished=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=1800, month=1, day=1, tzinfo=UTC), + finished=datetime(year=1900, month=1, day=1, tzinfo=UTC), archive_path=mnt_filepath.archive_path(enum_type=df_coll_type), ) @@ -45,12 +45,12 @@ class TestDFCollectionItemProperties: @pytest.mark.parametrize("df_coll_type", df_collection_types) class TestDFCollectionItemMethods: - def test_has_mysql_false(self, mnt_filepath: "GRLDatasets", df_coll_type): + def test_has_mysql_false(self, mnt_filepath: GRLDatasets, df_coll_type): collection = DFCollection( data_type=df_coll_type, offset="100d", - start=datetime(year=1800, month=1, day=1, tzinfo=timezone.utc), - finished=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=1800, month=1, day=1, tzinfo=UTC), + finished=datetime(year=1900, month=1, day=1, tzinfo=UTC), archive_path=mnt_filepath.archive_path(enum_type=df_coll_type), ) @@ -58,13 +58,13 @@ class TestDFCollectionItemMethods: assert not instance1.has_mysql() def test_has_mysql_true( - self, thl_web_rr: PostgresConfig, mnt_filepath: "GRLDatasets", df_coll_type + self, thl_web_rr: PostgresConfig, mnt_filepath: GRLDatasets, df_coll_type ): collection = DFCollection( data_type=df_coll_type, offset="100d", - start=datetime(year=1800, month=1, day=1, tzinfo=timezone.utc), - finished=datetime(year=1900, month=1, day=1, tzinfo=timezone.utc), + start=datetime(year=1800, month=1, day=1, tzinfo=UTC), + finished=datetime(year=1900, month=1, day=1, tzinfo=UTC), archive_path=mnt_filepath.archive_path(enum_type=df_coll_type), pg_config=thl_web_rr, ) diff --git a/tests/incite/collections/test_df_collection_item_thl_web.py b/tests/incite/collections/test_df_collection_item_thl_web.py index 8b8bcbe..062171d 100644 --- a/tests/incite/collections/test_df_collection_item_thl_web.py +++ b/tests/incite/collections/test_df_collection_item_thl_web.py @@ -1,11 +1,11 @@ from __future__ import annotations -from collections.abc import Generator -from datetime import datetime, timedelta, timezone +from collections.abc import Callable, Generator +from datetime import UTC, datetime, timedelta, timezone from itertools import product as iter_product from os.path import join as pjoin from pathlib import Path, PurePath -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 import dask.dataframe as dd @@ -52,7 +52,7 @@ unsupported_mock_types = { } -def combo_object() -> Generator[str, None, None]: +def combo_object() -> Generator[str]: for x in iter_product( df_collections, ["15min", "45min", "1H"], @@ -632,7 +632,7 @@ class TestDFCollectionItemMethodBase: for item in df_collection.items: item: DFCollectionItem - if datetime.now(tz=timezone.utc) > item.finish + aa: + if datetime.now(tz=UTC) > item.finish + aa: assert item.should_archive() else: assert not item.should_archive() diff --git a/tests/incite/collections/test_df_collection_thl_marketplaces.py b/tests/incite/collections/test_df_collection_thl_marketplaces.py index 981f62e..2597d38 100644 --- a/tests/incite/collections/test_df_collection_thl_marketplaces.py +++ b/tests/incite/collections/test_df_collection_thl_marketplaces.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from itertools import product from typing import TYPE_CHECKING @@ -57,8 +57,8 @@ class TestDFCollection_thl_marketplaces: archive_path=mnt_filepath.archive_path(enum_type=data_type), sql_helper=spectrum_rw, offset=offset, - start=datetime(year=2023, month=6, day=1, minute=0, tzinfo=timezone.utc), - finished=datetime(year=2023, month=6, day=1, minute=5, tzinfo=timezone.utc), + start=datetime(year=2023, month=6, day=1, minute=0, tzinfo=UTC), + finished=datetime(year=2023, month=6, day=1, minute=5, tzinfo=UTC), ) assert isinstance(instance, DFCollection) diff --git a/tests/incite/collections/test_df_collection_thl_web.py b/tests/incite/collections/test_df_collection_thl_web.py index b09d44c..2cb0ba0 100644 --- a/tests/incite/collections/test_df_collection_thl_web.py +++ b/tests/incite/collections/test_df_collection_thl_web.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: ) -def combo_object() -> Generator[tuple, None, None]: +def combo_object() -> Generator[tuple]: for x in product( [ DFCollectionType.USER, diff --git a/tests/incite/mergers/foundations/test_enriched_session.py b/tests/incite/mergers/foundations/test_enriched_session.py index 47f243e..a0ae01e 100644 --- a/tests/incite/mergers/foundations/test_enriched_session.py +++ b/tests/incite/mergers/foundations/test_enriched_session.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from itertools import product from typing import Optional @@ -77,15 +77,15 @@ class TestEnrichedSession: class TestEnrichedSessionAdmin: @pytest.fixture - def start(self) -> "datetime": - return datetime(year=2020, month=3, day=14, tzinfo=timezone.utc) + def start(self) -> datetime: + return datetime(year=2020, month=3, day=14, tzinfo=UTC) @pytest.fixture def offset(self) -> str: return "1d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return timedelta(days=5) def test_to_admin_response( diff --git a/tests/incite/mergers/foundations/test_enriched_wall.py b/tests/incite/mergers/foundations/test_enriched_wall.py index 8f4995b..b421df8 100644 --- a/tests/incite/mergers/foundations/test_enriched_wall.py +++ b/tests/incite/mergers/foundations/test_enriched_wall.py @@ -1,4 +1,4 @@ -from datetime import timedelta, timezone, datetime +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from itertools import product as iter_product from typing import Optional @@ -9,13 +9,13 @@ import pytest # noinspection PyUnresolvedReferences from distributed.utils_test import ( - gen_cluster, + cleanup, + client, client_no_amm, + cluster_fixture, + gen_cluster, loop, loop_in_thread, - cleanup, - cluster_fixture, - client, ) from generalresearch.incite.mergers.foundations.enriched_wall import ( @@ -158,15 +158,15 @@ class TestEnrichedWall: class TestEnrichedWallToAdmin: @pytest.fixture - def start(self) -> "datetime": - return datetime(year=2020, month=3, day=14, tzinfo=timezone.utc) + def start(self) -> datetime: + return datetime(year=2020, month=3, day=14, tzinfo=UTC) @pytest.fixture def offset(self) -> str: return "1d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return timedelta(days=5) def test_empty(self, enriched_wall_merge, client_no_amm, start): diff --git a/tests/incite/mergers/foundations/test_user_id_product.py b/tests/incite/mergers/foundations/test_user_id_product.py index f96bfb4..a696b45 100644 --- a/tests/incite/mergers/foundations/test_user_id_product.py +++ b/tests/incite/mergers/foundations/test_user_id_product.py @@ -1,4 +1,4 @@ -from datetime import timedelta, datetime, timezone +from datetime import UTC, datetime, timedelta, timezone from itertools import product import pandas as pd @@ -6,13 +6,13 @@ import pytest # noinspection PyUnresolvedReferences from distributed.utils_test import ( - gen_cluster, + cleanup, + client, client_no_amm, + cluster_fixture, + gen_cluster, loop, loop_in_thread, - cleanup, - cluster_fixture, - client, ) from generalresearch.incite.mergers.foundations.user_id_product import ( @@ -27,11 +27,7 @@ from test_utils.incite.mergers.conftest import user_id_product_merge product( ["12h", "3D"], [timedelta(days=5)], - [ - (datetime.now(tz=timezone.utc) - timedelta(days=35)).replace( - microsecond=0 - ) - ], + [(datetime.now(tz=UTC) - timedelta(days=35)).replace(microsecond=0)], ) ), ) diff --git a/tests/incite/mergers/test_merge_collection.py b/tests/incite/mergers/test_merge_collection.py index ec507bc..77fa8c7 100644 --- a/tests/incite/mergers/test_merge_collection.py +++ b/tests/incite/mergers/test_merge_collection.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from itertools import product import pandas as pd @@ -21,11 +21,7 @@ merge_types = list(e for e in MergeType if e != MergeType.TEST) merge_types, ["5min", "6h", "14D"], [timedelta(days=30)], - [ - (datetime.now(tz=timezone.utc) - timedelta(days=35)).replace( - microsecond=0 - ) - ], + [(datetime.now(tz=UTC) - timedelta(days=35)).replace(microsecond=0)], ) ), ) diff --git a/tests/incite/mergers/test_pop_ledger.py b/tests/incite/mergers/test_pop_ledger.py index 6f96108..7583faf 100644 --- a/tests/incite/mergers/test_pop_ledger.py +++ b/tests/incite/mergers/test_pop_ledger.py @@ -1,4 +1,4 @@ -from datetime import timedelta, datetime, timezone +from datetime import UTC, datetime, timedelta, timezone from itertools import product as iter_product from typing import Optional @@ -10,7 +10,7 @@ from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) from test_utils.incite.collections.conftest import ledger_collection -from test_utils.incite.conftest import mnt_filepath, incite_item_factory +from test_utils.incite.conftest import incite_item_factory, mnt_filepath from test_utils.incite.mergers.conftest import pop_ledger_merge from test_utils.managers.ledger.conftest import create_main_accounts @@ -27,11 +27,11 @@ from test_utils.managers.ledger.conftest import create_main_accounts class TestMergePOPLedger: @pytest.fixture - def start(self) -> "datetime": - return datetime(year=2020, month=3, day=14, tzinfo=timezone.utc) + def start(self) -> datetime: + return datetime(year=2020, month=3, day=14, tzinfo=UTC) @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return timedelta(days=5) def test_base( @@ -145,9 +145,9 @@ class TestMergePOPLedger: delete_ledger_db, session_collection, ): + from generalresearch.models.thl.finance import ProductBalances from generalresearch.models.thl.ledger import LedgerAccount from generalresearch.models.thl.product import Product - from generalresearch.models.thl.finance import ProductBalances u = user_factory(product=product, created=session_collection.start) diff --git a/tests/incite/mergers/test_ym_survey_merge.py b/tests/incite/mergers/test_ym_survey_merge.py index 4c2df6b..9107f21 100644 --- a/tests/incite/mergers/test_ym_survey_merge.py +++ b/tests/incite/mergers/test_ym_survey_merge.py @@ -1,4 +1,4 @@ -from datetime import timedelta, timezone, datetime +from datetime import UTC, datetime, timedelta, timezone from itertools import product import pandas as pd @@ -6,16 +6,16 @@ import pytest # noinspection PyUnresolvedReferences from distributed.utils_test import ( - gen_cluster, + cleanup, + client, client_no_amm, + cluster_fixture, + gen_cluster, loop, loop_in_thread, - cleanup, - cluster_fixture, - client, ) -from test_utils.incite.collections.conftest import wall_collection, session_collection +from test_utils.incite.collections.conftest import session_collection, wall_collection from test_utils.incite.mergers.conftest import ( enriched_session_merge, ym_survey_wall_merge, @@ -28,11 +28,7 @@ from test_utils.incite.mergers.conftest import ( product( ["12h", "3D"], [timedelta(days=30)], - [ - (datetime.now(tz=timezone.utc) - timedelta(days=35)).replace( - microsecond=0 - ) - ], + [(datetime.now(tz=UTC) - timedelta(days=35)).replace(microsecond=0)], ) ), ) diff --git a/tests/incite/schemas/test_admin_responses.py b/tests/incite/schemas/test_admin_responses.py index 43aa399..29d93fe 100644 --- a/tests/incite/schemas/test_admin_responses.py +++ b/tests/incite/schemas/test_admin_responses.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from random import sample from typing import List @@ -8,8 +8,8 @@ import pytest from generalresearch.incite.schemas import empty_dataframe_from_schema from generalresearch.incite.schemas.admin_responses import ( - AdminPOPSchema, SIX_HOUR_SECONDS, + AdminPOPSchema, ) from generalresearch.locales import Localelator @@ -72,8 +72,7 @@ class TestAdminPOPSchema: def test_index_tz_parser(self): tz_dates = [ - datetime(year=2024, month=1, day=i, tzinfo=timezone.utc) - for i in range(1, 10) + datetime(year=2024, month=1, day=i, tzinfo=UTC) for i in range(1, 10) ] df = pd.DataFrame( @@ -85,16 +84,16 @@ class TestAdminPOPSchema: df = self.assign_valid_vals(df) # Initially, they're all set with a timezone - timestmaps: List[pd.Timestamp] = [i for i in df.index.get_level_values(0)] - assert all([ts.tz == timezone.utc for ts in timestmaps]) + timestmaps: list[pd.Timestamp] = [i for i in df.index.get_level_values(0)] + assert all([ts.tz == UTC for ts in timestmaps]) # After validation, the timezone is removed df = AdminPOPSchema.validate(df) - timestmaps: List[pd.Timestamp] = [i for i in df.index.get_level_values(0)] + timestmaps: list[pd.Timestamp] = [i for i in df.index.get_level_values(0)] assert all([ts.tz is None for ts in timestmaps]) def test_index_tz_no_future_beyond_one_year(self): - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) tz_dates = [now + timedelta(days=i * 365) for i in range(1, 10)] df = pd.DataFrame( @@ -157,9 +156,7 @@ class TestAdminPOPSchema: def test_invalid_parsing(self): # (1) Timezones AND as strings will still parse correctly tz_str_dates = [ - datetime( - year=2024, month=1, day=1, minute=i, tzinfo=timezone.utc - ).isoformat() + datetime(year=2024, month=1, day=1, minute=i, tzinfo=UTC).isoformat() for i in range(1, 10) ] df = pd.DataFrame( @@ -173,12 +170,12 @@ class TestAdminPOPSchema: df = AdminPOPSchema.validate(df, lazy=True) assert isinstance(df, pd.DataFrame) - timestmaps: List[pd.Timestamp] = [i for i in df.index.get_level_values(0)] + timestmaps: list[pd.Timestamp] = [i for i in df.index.get_level_values(0)] assert all([ts.tz is None for ts in timestmaps]) # (2) Timezones are removed dates = [ - datetime(year=2024, month=1, day=1, minute=i, tzinfo=timezone.utc) + datetime(year=2024, month=1, day=1, minute=i, tzinfo=UTC) for i in range(1, 10) ] df = pd.DataFrame( @@ -190,12 +187,12 @@ class TestAdminPOPSchema: df = self.assign_valid_vals(df) # Has tz before validation, and none after - timestmaps: List[pd.Timestamp] = [i for i in df.index.get_level_values(0)] - assert all([ts.tz is timezone.utc for ts in timestmaps]) + timestmaps: list[pd.Timestamp] = [i for i in df.index.get_level_values(0)] + assert all([ts.tz is UTC for ts in timestmaps]) df = AdminPOPSchema.validate(df, lazy=True) - timestmaps: List[pd.Timestamp] = [i for i in df.index.get_level_values(0)] + timestmaps: list[pd.Timestamp] = [i for i in df.index.get_level_values(0)] assert all([ts.tz is None for ts in timestmaps]) def test_clipping(self): diff --git a/tests/incite/test_collection_base.py b/tests/incite/test_collection_base.py index 7e6605f..5a63019 100644 --- a/tests/incite/test_collection_base.py +++ b/tests/incite/test_collection_base.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from os.path import exists as pexists from os.path import join as pjoin from pathlib import Path @@ -12,11 +12,9 @@ from _pytest._code.code import ExceptionInfo from generalresearch.incite.base import CollectionBase from test_utils.incite.conftest import mnt_filepath -AGO_15min = (datetime.now(tz=timezone.utc) - timedelta(minutes=15)).replace( - microsecond=0 -) -AGO_1HR = (datetime.now(tz=timezone.utc) - timedelta(hours=1)).replace(microsecond=0) -AGO_2HR = (datetime.now(tz=timezone.utc) - timedelta(hours=2)).replace(microsecond=0) +AGO_15min = (datetime.now(tz=UTC) - timedelta(minutes=15)).replace(microsecond=0) +AGO_1HR = (datetime.now(tz=UTC) - timedelta(hours=1)).replace(microsecond=0) +AGO_2HR = (datetime.now(tz=UTC) - timedelta(hours=2)).replace(microsecond=0) class TestCollectionBase: @@ -50,7 +48,7 @@ class TestCollectionBase: with pytest.raises(expected_exception=ValueError) as cm: cm: ExceptionInfo CollectionBase( - start=datetime.now(tz=timezone.utc) - timedelta(days=10), + start=datetime.now(tz=UTC) - timedelta(days=10), archive_path=mnt_filepath.data_src, ) assert "Collection.start must not have microseconds" in str(cm.value) @@ -66,9 +64,7 @@ class TestCollectionBase: assert "Timezone is not UTC" in str(cm.value) instance = CollectionBase(archive_path=mnt_filepath.data_src) - assert instance.start == datetime( - year=2018, month=1, day=1, tzinfo=timezone.utc - ) + assert instance.start == datetime(year=2018, month=1, day=1, tzinfo=UTC) with pytest.raises(expected_exception=ValueError) as cm: cm: ExceptionInfo @@ -145,7 +141,7 @@ class TestCollectionBaseProperties: instance._interval_range(end=datetime.now(tz=tz)) assert "Timezones must match" in str(cm.value) - res = instance._interval_range(end=datetime.now(tz=timezone.utc)) + res = instance._interval_range(end=datetime.now(tz=UTC)) assert isinstance(res, pd.IntervalIndex) assert res.closed_left assert res.is_non_overlapping_monotonic @@ -282,7 +278,7 @@ class TestCollectionBaseMethodsSourceTiming: def test_get_item_start(self, mnt_filepath): instance = CollectionBase(archive_path=mnt_filepath.data_src) - dt = datetime.now(tz=timezone.utc) + dt = datetime.now(tz=UTC) start = pd.Timestamp(dt) with pytest.raises(expected_exception=NotImplementedError) as cm: @@ -292,7 +288,7 @@ class TestCollectionBaseMethodsSourceTiming: def test_get_items(self, mnt_filepath): instance = CollectionBase(archive_path=mnt_filepath.data_src) - dt = datetime.now(tz=timezone.utc) + dt = datetime.now(tz=UTC) with pytest.raises(expected_exception=NotImplementedError) as cm: instance.get_items(since=dt) diff --git a/tests/incite/test_collection_base_item.py b/tests/incite/test_collection_base_item.py index e5d1d02..3f4d023 100644 --- a/tests/incite/test_collection_base_item.py +++ b/tests/incite/test_collection_base_item.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from os.path import join as pjoin from pathlib import Path from uuid import uuid4 @@ -13,7 +13,7 @@ from generalresearch.incite.base import CollectionItemBase class TestCollectionItemBase: def test_init(self): - dt = datetime.now(tz=timezone.utc).replace(microsecond=0) + dt = datetime.now(tz=UTC).replace(microsecond=0) instance = CollectionItemBase() instance2 = CollectionItemBase(start=dt) @@ -25,7 +25,7 @@ class TestCollectionItemBase: assert 0 == instance.start.microsecond == instance2.start.microsecond def test_init_start(self): - dt = datetime.now(tz=timezone.utc) + dt = datetime.now(tz=UTC) with pytest.raises(expected_exception=ValidationError) as cm: CollectionItemBase(start=dt) diff --git a/tests/managers/leaderboard.py b/tests/managers/leaderboard.py index 4d32dd0..149bdbb 100644 --- a/tests/managers/leaderboard.py +++ b/tests/managers/leaderboard.py @@ -1,7 +1,7 @@ import os import time import zoneinfo -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal from uuid import uuid4 @@ -10,9 +10,6 @@ import pytest from generalresearch.managers.leaderboard.manager import LeaderboardManager from generalresearch.managers.leaderboard.tasks import hit_leaderboards from generalresearch.models.thl.definitions import Status -from generalresearch.models.thl.user import User -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session from generalresearch.models.thl.leaderboard import ( LeaderboardCode, LeaderboardFrequency, @@ -22,7 +19,10 @@ from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, PayoutTransformationPercentArgs, + Product, ) +from generalresearch.models.thl.session import Session +from generalresearch.models.thl.user import User # random uuid for leaderboard tests product_id = uuid4().hex @@ -63,7 +63,7 @@ def _create_session( ) session = Session( user=user, - started=datetime(2025, 2, 5, 6, tzinfo=timezone.utc), + started=datetime(2025, 2, 5, 6, tzinfo=UTC), id=1, country_iso=country_iso, status=Status.COMPLETE, @@ -152,7 +152,7 @@ class TestLeaderboards: 999999, tzinfo=zoneinfo.ZoneInfo(key="America/New_York"), ) - assert lb.period_start_utc == datetime(2025, 2, 5, 5, tzinfo=timezone.utc) + assert lb.period_start_utc == datetime(2025, 2, 5, 5, tzinfo=UTC) assert lb.row_count == 7 assert lb.rows == [ LeaderboardRow(bpuid="aaa", rank=1, value=10), @@ -270,5 +270,5 @@ class TestLeaderboards: ) assert lb.local_start_time == "2025-02-01T00:00:00+09:00" assert lb.local_end_time == "2025-02-01T23:59:59.999999+09:00" - assert lb.period_start_utc == datetime(2025, 1, 31, 15, tzinfo=timezone.utc) + assert lb.period_start_utc == datetime(2025, 1, 31, 15, tzinfo=UTC) print(lb.model_dump(mode="json")) diff --git a/tests/managers/test_events.py b/tests/managers/test_events.py index a0fab38..6941c00 100644 --- a/tests/managers/test_events.py +++ b/tests/managers/test_events.py @@ -1,22 +1,22 @@ +import math import random import time -from datetime import timedelta, datetime, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from functools import partial +from math import floor from typing import Optional from uuid import uuid4 -import math import pytest -from math import floor from generalresearch.managers.events import EventSubscriber from generalresearch.models import Source from generalresearch.models.events import ( - MessageKind, - EventType, AggregateBySource, + EventType, MaxGaugeBySource, + MessageKind, ) from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.definitions import Status, StatusCode1 @@ -41,13 +41,13 @@ def event_subscriber(thl_redis_config, product_id): def create_dummy( - product_id: Optional[str] = None, product_user_id: Optional[str] = None + product_id: str | None = None, product_user_id: str | None = None ) -> User: return User( product_id=product_id, product_user_id=product_user_id or uuid4().hex, uuid=uuid4().hex, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), user_id=random.randint(0, floor(2**32 / 2)), ) @@ -496,7 +496,7 @@ class TestChannelsSubscriptions: wall.update( status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - finished=datetime.now(tz=timezone.utc), + finished=datetime.now(tz=UTC), cpi=Decimal("1"), ) event_manager.handle_task_finish(wall, session, user) diff --git a/tests/managers/thl/test_contest/test_leaderboard.py b/tests/managers/thl/test_contest/test_leaderboard.py index 80a88a5..1a52f83 100644 --- a/tests/managers/thl/test_contest/test_leaderboard.py +++ b/tests/managers/thl/test_contest/test_leaderboard.py @@ -1,10 +1,10 @@ -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from zoneinfo import ZoneInfo from generalresearch.currency import USDCent from generalresearch.models.thl.contest.definitions import ( - ContestStatus, ContestEndReason, + ContestStatus, ) from generalresearch.models.thl.contest.leaderboard import ( LeaderboardContest, @@ -13,9 +13,11 @@ from generalresearch.models.thl.contest.leaderboard import ( from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User from test_utils.managers.contest.conftest import ( - leaderboard_contest_in_db as contest_in_db, leaderboard_contest_create as contest_create, ) +from test_utils.managers.contest.conftest import ( + leaderboard_contest_in_db as contest_in_db, +) class TestLeaderboardContestCRUD: @@ -39,7 +41,7 @@ class TestLeaderboardContestCRUD: # We have it set in the fixture as the daily contest for 2025-01-01 assert c.end_condition.ends_at == datetime( 2025, 1, 1, 23, 59, 59, 999999, tzinfo=ZoneInfo("America/New_York") - ).astimezone(tz=timezone.utc) + timedelta(minutes=90) + ).astimezone(tz=UTC) + timedelta(minutes=90) def test_enter( self, diff --git a/tests/managers/thl/test_contest/test_milestone.py b/tests/managers/thl/test_contest/test_milestone.py index 7312a64..66c5dc4 100644 --- a/tests/managers/thl/test_contest/test_milestone.py +++ b/tests/managers/thl/test_contest/test_milestone.py @@ -1,23 +1,29 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from generalresearch.models.thl.contest.definitions import ( - ContestStatus, ContestEndReason, + ContestStatus, ) from generalresearch.models.thl.contest.milestone import ( + ContestEntryTrigger, MilestoneContest, MilestoneContestCreate, MilestoneUserView, - ContestEntryTrigger, ) from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User from test_utils.managers.contest.conftest import ( milestone_contest as contest, - milestone_contest_in_db as contest_in_db, +) +from test_utils.managers.contest.conftest import ( milestone_contest_create as contest_create, +) +from test_utils.managers.contest.conftest import ( milestone_contest_factory as contest_factory, ) +from test_utils.managers.contest.conftest import ( + milestone_contest_in_db as contest_in_db, +) class TestMilestoneContest: @@ -28,7 +34,7 @@ class TestMilestoneContest: assert not should, msg # Change so that the contest ends now - contest.end_condition.ends_at = datetime.now(tz=timezone.utc) + contest.end_condition.ends_at = datetime.now(tz=UTC) should, msg = contest.should_end() assert should assert msg == ContestEndReason.ENDS_AT diff --git a/tests/managers/thl/test_contest/test_raffle.py b/tests/managers/thl/test_contest/test_raffle.py index 060055a..5804ea3 100644 --- a/tests/managers/thl/test_contest/test_raffle.py +++ b/tests/managers/thl/test_contest/test_raffle.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone import pytest from pydantic import ValidationError @@ -9,21 +9,19 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionConditionFailedError, ) from generalresearch.models.thl.contest import ( - ContestPrize, - ContestEntryRule, ContestEndCondition, + ContestEntryRule, + ContestPrize, ) from generalresearch.models.thl.contest.definitions import ( - ContestStatus, - ContestPrizeKind, ContestEndReason, + ContestPrizeKind, + ContestStatus, ) from generalresearch.models.thl.contest.exceptions import ContestError from generalresearch.models.thl.contest.raffle import ( ContestEntry, ContestEntryType, -) -from generalresearch.models.thl.contest.raffle import ( RaffleContest, RaffleContestCreate, RaffleUserView, @@ -32,10 +30,16 @@ from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User from test_utils.managers.contest.conftest import ( raffle_contest as contest, - raffle_contest_in_db as contest_in_db, +) +from test_utils.managers.contest.conftest import ( raffle_contest_create as contest_create, +) +from test_utils.managers.contest.conftest import ( raffle_contest_factory as contest_factory, ) +from test_utils.managers.contest.conftest import ( + raffle_contest_in_db as contest_in_db, +) class TestRaffleContest: @@ -46,7 +50,7 @@ class TestRaffleContest: assert not should, msg # Change so that the contest ends now - contest.end_condition.ends_at = datetime.now(tz=timezone.utc) + contest.end_condition.ends_at = datetime.now(tz=UTC) should, msg = contest.should_end() assert should assert msg == ContestEndReason.ENDS_AT diff --git a/tests/managers/thl/test_harmonized_uqa.py b/tests/managers/thl/test_harmonized_uqa.py index 6bbbbe1..3b6df48 100644 --- a/tests/managers/thl/test_harmonized_uqa.py +++ b/tests/managers/thl/test_harmonized_uqa.py @@ -1,11 +1,11 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone import pytest from generalresearch.managers.thl.profiling.uqa import UQAManager from generalresearch.models.thl.profiling.user_question_answer import ( - UserQuestionAnswer, DUMMY_UQA, + UserQuestionAnswer, ) from generalresearch.models.thl.user import User @@ -18,7 +18,7 @@ class TestUQAManager: assert len(uqas) == 0 def test_create(self, uqa_manager: UQAManager, user: User): - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) uqas = [ UserQuestionAnswer( user_id=user.user_id, @@ -38,7 +38,7 @@ class TestUQAManager: assert res[0] == uqas[0] # Same question, so this gets updated - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) uqas_update = [ UserQuestionAnswer( user_id=user.user_id, @@ -57,7 +57,7 @@ class TestUQAManager: assert res[0] == uqas_update[0] # Add a new answer - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) uqas_new = [ UserQuestionAnswer( user_id=user.user_id, @@ -103,7 +103,7 @@ class TestUQAManagerCache: UserQuestionAnswer( question_id="5d6d9f3c03bb40bf9d0a24f306387d7c", answer=("1",), - timestamp=datetime.now(tz=timezone.utc), + timestamp=datetime.now(tz=UTC), country_iso="us", language_iso="eng", property_code="gr:gender", diff --git a/tests/managers/thl/test_ledger/test_lm_accounts.py b/tests/managers/thl/test_ledger/test_lm_accounts.py index 5cfaac1..faef5fb 100644 --- a/tests/managers/thl/test_ledger/test_lm_accounts.py +++ b/tests/managers/thl/test_ledger/test_lm_accounts.py @@ -1,6 +1,7 @@ +from collections.abc import Callable from itertools import product as iproduct from random import randint -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -51,10 +52,10 @@ class TestLedgerAccountManagerNoResults: def test_get_account_no_results( self, - currency: "LedgerCurrency", + currency: LedgerCurrency, kind: str, - acct_id: "UUIDStr", - lm: "LedgerManager", + acct_id: UUIDStr, + lm: LedgerManager, ): """Try to query for accounts that we know don't exist and confirm that we either get the expected None result or it raises the correct @@ -74,10 +75,10 @@ class TestLedgerAccountManagerNoResults: def test_get_account_no_results_many( self, - currency: "LedgerCurrency", + currency: LedgerCurrency, kind: str, - acct_id: "UUIDStr", - lm: "LedgerManager", + acct_id: UUIDStr, + lm: LedgerManager, ): qn = ":".join([currency, kind, acct_id]) @@ -114,10 +115,10 @@ class TestLedgerAccountManagerCreate: def test_create_account_error_permission( self, - currency: "LedgerCurrency", - account_type: "AccountType", - direction: "Direction", - lm: "LedgerManager", + currency: LedgerCurrency, + account_type: AccountType, + direction: Direction, + lm: LedgerManager, ): """Confirm that the Permission values that are set on the Ledger Manger allow the Creation action to occur. @@ -164,10 +165,10 @@ class TestLedgerAccountManagerCreate: def test_create( self, - currency: "LedgerCurrency", - account_type: "AccountType", - direction: "Direction", - lm: "LedgerManager", + currency: LedgerCurrency, + account_type: AccountType, + direction: Direction, + lm: LedgerManager, ): """Confirm that the Permission values that are set on the Ledger Manger allow the Creation action to occur. @@ -194,10 +195,10 @@ class TestLedgerAccountManagerCreate: def test_get_or_create( self, - currency: "LedgerCurrency", - account_type: "AccountType", - direction: "Direction", - lm: "LedgerManager", + currency: LedgerCurrency, + account_type: AccountType, + direction: Direction, + lm: LedgerManager, ): """Confirm that the Permission values that are set on the Ledger Manger allow the Creation action to occur. @@ -225,7 +226,7 @@ class TestLedgerAccountManagerCreate: class TestLedgerAccountManagerGet: - def test_get(self, ledger_account: "LedgerAccount", lm: "LedgerManager"): + def test_get(self, ledger_account: LedgerAccount, lm: LedgerManager): res = lm.get_account(qualified_name=ledger_account.qualified_name) assert res is not None assert res.uuid == ledger_account.uuid @@ -243,11 +244,11 @@ class TestLedgerAccountManagerGet: def test_get_balance_empty( self, - ledger_account: "LedgerAccount", - ledger_account_credit: "LedgerAccount", - ledger_account_debit: "LedgerAccount", - ledger_tx: "LedgerTransaction", - lm: "LedgerManager", + ledger_account: LedgerAccount, + ledger_account_credit: LedgerAccount, + ledger_account_debit: LedgerAccount, + ledger_tx: LedgerTransaction, + lm: LedgerManager, ): res = lm.get_account_balance(account=ledger_account) assert res == 0 @@ -261,12 +262,12 @@ class TestLedgerAccountManagerGet: @pytest.mark.parametrize("n_times", range(5)) def test_get_account_filtered_balance( self, - ledger_account: "LedgerAccount", - ledger_account_credit: "LedgerAccount", - ledger_account_debit: "LedgerAccount", - ledger_tx: "LedgerTransaction", - n_times: "PositiveInt", - lm: "LedgerManager", + ledger_account: LedgerAccount, + ledger_account_credit: LedgerAccount, + ledger_account_debit: LedgerAccount, + ledger_tx: LedgerTransaction, + n_times: PositiveInt, + lm: LedgerManager, ): """Try searching for random metadata and confirm it's always 0 because Tx can be found. @@ -320,7 +321,7 @@ class TestLedgerAccountManagerGet: ) def test_get_balance_timerange_empty( - self, ledger_account: "LedgerAccount", lm: "LedgerManager" + self, ledger_account: LedgerAccount, lm: LedgerManager ): res = lm.get_account_balance_timerange(account=ledger_account) assert res == 0 diff --git a/tests/managers/thl/test_ledger/test_lm_tx_locks.py b/tests/managers/thl/test_ledger/test_lm_tx_locks.py index df2611b..07c3712 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx_locks.py +++ b/tests/managers/thl/test_ledger/test_lm_tx_locks.py @@ -1,7 +1,7 @@ import logging -from datetime import datetime, timezone, timedelta +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import Callable import pytest @@ -9,21 +9,21 @@ from generalresearch.managers.thl.ledger_manager.conditions import ( generate_condition_mp_payment, ) from generalresearch.managers.thl.ledger_manager.exceptions import ( + LedgerTransactionCreateError, LedgerTransactionCreateLockError, LedgerTransactionFlagAlreadyExistsError, - LedgerTransactionCreateError, ) from generalresearch.models import Source from generalresearch.models.thl.ledger import LedgerTransaction from generalresearch.models.thl.session import ( - Wall, + Session, Status, StatusCode1, - Session, + Wall, WallAdjustedStatus, ) from generalresearch.models.thl.user import User -from test_utils.models.conftest import user_factory, session, product_user_wallet_no +from test_utils.models.conftest import product_user_wallet_no, session, user_factory logger = logging.getLogger("LedgerManager") @@ -139,7 +139,7 @@ class TestLedgerLocks: delete_ledger_db() create_main_accounts() - now = datetime.now(timezone.utc) - timedelta(hours=1) + now = datetime.now(UTC) - timedelta(hours=1) user: User = user_factory(product=product_user_wallet_no) # A User does a Wall complete on Session.id=1 and the transaction is @@ -283,8 +283,8 @@ class TestLedgerLocks: session_id=3, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall=wall1, user=user, created=wall1.started) @@ -327,8 +327,8 @@ class TestLedgerLocks: session_id=3, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall1, user, created=wall1.started) diff --git a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py index 294d092..1fb9c01 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py @@ -1,5 +1,5 @@ import logging -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from random import randint from uuid import uuid4 @@ -11,22 +11,22 @@ from redis.lock import Lock from generalresearch.currency import USDCent from generalresearch.managers.base import Permission -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.managers.thl.ledger_manager.exceptions import ( - LedgerTransactionFlagAlreadyExistsError, LedgerTransactionConditionFailedError, - LedgerTransactionReleaseLockError, LedgerTransactionCreateError, + LedgerTransactionFlagAlreadyExistsError, + LedgerTransactionReleaseLockError, ) from generalresearch.managers.thl.ledger_manager.ledger import LedgerTransaction +from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models import Source from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.ledger import Direction, TransactionType from generalresearch.models.thl.session import ( - Wall, + Session, Status, StatusCode1, - Session, + Wall, ) from generalresearch.models.thl.user import User from generalresearch.models.thl.wallet import PayoutType @@ -55,7 +55,7 @@ class TestThlLedgerManagerBPPayout: delete_ledger_db() create_main_accounts() - now = datetime.now(timezone.utc) - timedelta(hours=1) + now = datetime.now(UTC) - timedelta(hours=1) user: User = user_factory(product=product_user_wallet_no) wall1 = Wall( @@ -158,7 +158,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_wallet_balance_check=True, skip_one_per_day_check=True, skip_flag_check=True, @@ -189,7 +189,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=uuid4().hex, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_wallet_balance_check=False, skip_one_per_day_check=False, skip_flag_check=False, @@ -199,7 +199,7 @@ class TestThlLedgerManagerBPPayout: def test_create_tx_redis_failure(self, product, thl_web_rw, thl_lm): rand_amount: USDCent = USDCent(randint(100, 1_000)) payoutevent_uuid = uuid4().hex - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) thl_lm.create_tx_plug_bp_wallet( product, rand_amount, now, direction=Direction.CREDIT @@ -226,7 +226,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) assert e.type is redis.exceptions.TimeoutError # No txs were created @@ -238,7 +238,7 @@ class TestThlLedgerManagerBPPayout: def test_create_tx_multiple_per_day(self, product, thl_lm): rand_amount: USDCent = USDCent(randint(100, 1_000)) payoutevent_uuid = uuid4().hex - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) thl_lm.create_tx_plug_bp_wallet( product, rand_amount * USDCent(2), now, direction=Direction.CREDIT @@ -248,7 +248,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) # Try to create another @@ -258,7 +258,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) assert e.type is LedgerTransactionFlagAlreadyExistsError @@ -270,7 +270,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid2, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) assert e.type is LedgerTransactionConditionFailedError assert str(e.value) == ">1 tx per day" @@ -280,14 +280,14 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid2, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_one_per_day_check=True, ) def test_create_tx_redis_lock_release_error(self, product, thl_lm): rand_amount: USDCent = USDCent(randint(100, 1_000)) payoutevent_uuid = uuid4().hex - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) bp_wallet_account = thl_lm.get_account_or_create_bp_wallet(product=product) thl_lm.create_tx_plug_bp_wallet( @@ -304,7 +304,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) assert e.type is LedgerTransactionCreateError assert str(e.value) == "Redis error: Simulated timeout during acquire" @@ -321,7 +321,7 @@ class TestThlLedgerManagerBPPayout: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) assert e.type is LedgerTransactionReleaseLockError assert str(e.value) == "Redis error: Simulated timeout during release" @@ -337,7 +337,7 @@ class TestPayoutEventManagerBPPayout: def test_create(self, product, thl_lm, brokerage_product_payout_event_manager): rand_amount: USDCent = USDCent(randint(100, 1_000)) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) bp_wallet_account = thl_lm.get_account_or_create_bp_wallet(product=product) assert thl_lm.get_account_balance(bp_wallet_account) == 0 thl_lm.create_tx_plug_bp_wallet( @@ -369,7 +369,7 @@ class TestPayoutEventManagerBPPayout: original_release = Lock.release rand_amount: USDCent = USDCent(randint(100, 1_000)) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) bp_wallet_account = thl_lm.get_account_or_create_bp_wallet(product=product) assert thl_lm.get_account_balance(bp_wallet_account) == 0 thl_lm.create_tx_plug_bp_wallet( @@ -435,7 +435,7 @@ class TestPayoutEventManagerBPPayout: # We wouldn't do this in practice, because this is paying out the BP again, but # we can if want to. # Change the timestamp so it'll create a new payout event - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) with pytest.raises(LedgerTransactionConditionFailedError) as e: pe = brokerage_product_payout_event_manager.create_bp_payout_event( thl_ledger_manager=thl_lm, @@ -450,7 +450,7 @@ class TestPayoutEventManagerBPPayout: assert pe.status == PayoutStatus.FAILED # And if we really want to, we can make it again - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) pe = brokerage_product_payout_event_manager.create_bp_payout_event( thl_ledger_manager=thl_lm, product=product, @@ -478,7 +478,7 @@ class TestPayoutEventManagerBPPayout: original_release = Lock.release rand_amount: USDCent = USDCent(randint(100, 1_000)) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) bp_wallet_account = thl_lm.get_account_or_create_bp_wallet(product=product) brokerage_product_payout_event_manager.set_account_lookup_table(thl_lm=thl_lm) diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx.py b/tests/managers/thl/test_ledger/test_thl_lm_tx.py index 31c7107..be988a1 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx.py @@ -1,5 +1,5 @@ import logging -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from random import randint from uuid import uuid4 @@ -14,23 +14,22 @@ from generalresearch.models import Source from generalresearch.models.thl.definitions import ( WALL_ALLOWED_STATUS_STATUS_CODE, ) -from generalresearch.models.thl.ledger import Direction -from generalresearch.models.thl.ledger import TransactionType +from generalresearch.models.thl.ledger import Direction, TransactionType +from generalresearch.models.thl.payout import UserPayoutEvent from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, UserWalletConfig, ) from generalresearch.models.thl.session import ( - Wall, + Session, Status, StatusCode1, - Session, + Wall, WallAdjustedStatus, ) from generalresearch.models.thl.user import User from generalresearch.models.thl.wallet import PayoutType -from generalresearch.models.thl.payout import UserPayoutEvent logger = logging.getLogger("LedgerManager") @@ -82,7 +81,7 @@ class TestThlLedgerTxManager: session=s1, status=Status.COMPLETE, status_code_1=status_code_1, - finished=datetime.now(tz=timezone.utc) + timedelta(minutes=10), + finished=datetime.now(tz=UTC) + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay, ) @@ -127,7 +126,7 @@ class TestThlLedgerTxManager: session=s1, status=Status.COMPLETE, status_code_1=status_code_1, - finished=datetime.now(tz=timezone.utc) + timedelta(minutes=10), + finished=datetime.now(tz=UTC) + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay, ) @@ -209,7 +208,7 @@ class TestThlLedgerTxManager: # there is no financial changes needed session.update( **{ - "finished": datetime.now(tz=timezone.utc) + timedelta(minutes=10), + "finished": datetime.now(tz=UTC) + timedelta(minutes=10), } ) assert session.finished @@ -229,7 +228,7 @@ class TestThlLedgerTxManager: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_wallet_balance_check=True, skip_one_per_day_check=True, skip_flag_check=True, @@ -260,7 +259,7 @@ class TestThlLedgerTxManager: product=product, amount=rand_amount, payoutevent_uuid=uuid4().hex, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_wallet_balance_check=False, skip_one_per_day_check=False, skip_flag_check=False, @@ -276,7 +275,7 @@ class TestThlLedgerTxManager: product=product, amount=rand_amount, payoutevent_uuid=payoutevent_uuid, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), ) # Check the basic attributes @@ -300,7 +299,7 @@ class TestThlLedgerTxManager: tx = thl_lm.create_tx_plug_bp_wallet( product=product, amount=rand_amount, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), direction=Direction.DEBIT, skip_flag_check=False, ) @@ -328,7 +327,7 @@ class TestThlLedgerTxManager: tx = thl_lm.create_tx_plug_bp_wallet_( product=product, amount=rand_amount, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), direction=Direction.DEBIT, ) @@ -345,7 +344,7 @@ class TestThlLedgerTxManager: thl_lm.create_tx_plug_bp_wallet_( product=product, amount=rand_amount + rand_amount, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), direction=Direction.CREDIT, ) balance = thl_lm.get_account_balance( @@ -727,8 +726,8 @@ class TestThlLedgerTxManagerFlows: session_id=1, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall=wall1, user=user, created=wall1.started) @@ -740,8 +739,8 @@ class TestThlLedgerTxManagerFlows: session_id=1, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall=wall2, user=user, created=wall2.started) @@ -793,8 +792,8 @@ class TestThlLedgerTxManagerFlows: session_id=1, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) tx = thl_lm.create_tx_task_complete( wall=wall1, user=user, created=wall1.started @@ -880,8 +879,8 @@ class TestThlLedgerTxManagerFlows: session_id=3, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) tx = thl_lm.create_tx_task_complete( @@ -922,8 +921,8 @@ class TestThlLedgerTxManagerFlows: session_id=3, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall=wall1, user=user, created=wall1.started) @@ -963,8 +962,8 @@ class TestThlLedgerTxManagerFlows: session_id=3, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) thl_lm.create_tx_task_complete(wall=wall1, user=user, created=wall1.started) @@ -1416,7 +1415,7 @@ class TestThlLedgerManagerAdj: delete_ledger_db() create_main_accounts() - now = datetime.now(timezone.utc) - timedelta(days=1) + now = datetime.now(UTC) - timedelta(days=1) user: User = user_factory(product=product_user_wallet_yes) # Create 2 Wall completes and create the respective transaction for diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py index 1e7146a..9253ff0 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py @@ -1,17 +1,17 @@ import logging -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from uuid import uuid4 import pytest from generalresearch.managers.thl.ledger_manager.exceptions import ( - LedgerTransactionFlagAlreadyExistsError, LedgerTransactionConditionFailedError, + LedgerTransactionFlagAlreadyExistsError, ) +from generalresearch.models.thl.payout import UserPayoutEvent from generalresearch.models.thl.user import User from generalresearch.models.thl.wallet import PayoutType -from generalresearch.models.thl.payout import UserPayoutEvent from test_utils.managers.ledger.conftest import create_main_accounts @@ -243,7 +243,7 @@ class TestLedgerManagerAMT: delete_ledger_db() create_main_accounts() - now = datetime.now(timezone.utc) - timedelta(hours=1) + now = datetime.now(UTC) - timedelta(hours=1) user: User = user_factory(product=product_amt_true) pe = UserPayoutEvent( @@ -394,7 +394,7 @@ class TestLedgerManagerPaypal: delete_ledger_db() create_main_accounts() - now = datetime.now(tz=timezone.utc) - timedelta(hours=1) + now = datetime.now(tz=UTC) - timedelta(hours=1) user: User = user_factory(product=product_amt_true) # debit_account_uuid nothing checks they match the ledger ... todo? diff --git a/tests/managers/thl/test_ledger/test_user_txs.py b/tests/managers/thl/test_ledger/test_user_txs.py index ecf146f..b4b0437 100644 --- a/tests/managers/thl/test_ledger/test_user_txs.py +++ b/tests/managers/thl/test_ledger/test_user_txs.py @@ -1,6 +1,7 @@ -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING from uuid import uuid4 from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager @@ -24,8 +25,8 @@ if TYPE_CHECKING: def test_user_txs( - user_factory: Callable[..., "User"], - product_amt_true: "Product", + user_factory: Callable[..., User], + product_amt_true: Product, create_main_accounts: Callable[..., None], thl_lm: ThlLedgerManager, lm, @@ -36,7 +37,7 @@ def test_user_txs( session_factory, user_payout_event_manager, utc_now: datetime, - settings: "GRLSettings", + settings: GRLSettings, ): delete_ledger_db() create_main_accounts() @@ -136,13 +137,13 @@ def test_user_txs( def test_user_txs_pagination( - user_factory: Callable[..., "User"], - product_amt_true: "Product", + user_factory: Callable[..., User], + product_amt_true: Product, create_main_accounts: Callable[..., None], - thl_lm: "ThlLedgerManager", - lm: "LedgerManager", + thl_lm: ThlLedgerManager, + lm: LedgerManager, delete_ledger_db: Callable[..., None], - session_with_tx_factory: Callable[..., "Session"], + session_with_tx_factory: Callable[..., Session], adj_to_fail_with_tx_factory, user_payout_event_manager, utc_now: datetime, @@ -187,7 +188,7 @@ def test_user_txs_pagination( assert txs.summary.user_bonus.entry_count == 12 # Test filtering. We should pull back only this one - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) user_compensate( ledger_manager=thl_lm, user=user, @@ -203,7 +204,7 @@ def test_user_txs_pagination( assert txs.summary.user_bonus.entry_count == 1 # And filtering with 0 results - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) txs = thl_lm.get_user_txs(user, page=1, size=5, time_start=now) assert len(txs.transactions) == 0 assert txs.total == 0 @@ -215,8 +216,8 @@ def test_user_txs_pagination( def test_user_txs_rolling_balance( - user_factory: Callable[..., "User"], - product_amt_true: "Product", + user_factory: Callable[..., User], + product_amt_true: Product, create_main_accounts, thl_lm, lm, @@ -224,7 +225,7 @@ def test_user_txs_rolling_balance( session_with_tx_factory, adj_to_fail_with_tx_factory, user_payout_event_manager, - settings: "GRLSettings", + settings: GRLSettings, ): """ Creates 3 $1.00 bonuses (postive), diff --git a/tests/managers/thl/test_maxmind.py b/tests/managers/thl/test_maxmind.py index c588c58..75bf0e9 100644 --- a/tests/managers/thl/test_maxmind.py +++ b/tests/managers/thl/test_maxmind.py @@ -1,17 +1,12 @@ import json import logging -from typing import Callable +from collections.abc import Callable -import geoip2.models import pytest from faker import Faker from faker.providers.address.en_US import Provider as USAddressProvider from generalresearch.managers.thl.ipinfo import GeoIpInfoManager -from generalresearch.managers.thl.maxmind import MaxmindManager -from generalresearch.managers.thl.maxmind.basic import ( - MaxmindBasicManager, -) from generalresearch.models.thl.ipinfo import ( GeoIPInformation, normalize_ip, diff --git a/tests/managers/thl/test_profiling/test_user_upk.py b/tests/managers/thl/test_profiling/test_user_upk.py index 53bb8fe..491e2b1 100644 --- a/tests/managers/thl/test_profiling/test_user_upk.py +++ b/tests/managers/thl/test_profiling/test_user_upk.py @@ -1,8 +1,8 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from generalresearch.managers.thl.profiling.user_upk import UserUpkManager -now = datetime.now(tz=timezone.utc) +now = datetime.now(tz=UTC) base = { "country_iso": "us", "language_iso": "eng", diff --git a/tests/managers/thl/test_survey.py b/tests/managers/thl/test_survey.py index 58c4577..4b4a579 100644 --- a/tests/managers/thl/test_survey.py +++ b/tests/managers/thl/test_survey.py @@ -1,24 +1,24 @@ import uuid -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal import pytest from generalresearch.models import Source from generalresearch.models.legacy.bucket import ( - SurveyEligibilityCriterion, - TopNPlusBucket, DurationSummary, PayoutSummary, + SurveyEligibilityCriterion, + TopNPlusBucket, ) from generalresearch.models.thl.profiling.user_question_answer import ( UserQuestionAnswer, ) from generalresearch.models.thl.survey.model import ( Survey, - SurveyStat, SurveyCategoryModel, SurveyEligibilityDefinition, + SurveyStat, ) @@ -258,7 +258,7 @@ class TestSurveyStat: return # 1,000 of the 20,000 are "new" - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) for s in ss[:1000]: s.survey__survey_id = "b" s.updated_at = now @@ -298,7 +298,7 @@ class TestSurveyStat: source=source, surveys=surveys, survey_stats=survey_stats ) # UPDATE ------- - since = datetime.now(tz=timezone.utc) + since = datetime.now(tz=UTC) print(f"{since=}") # 10 survey disappear diff --git a/tests/managers/thl/test_task_adjustment.py b/tests/managers/thl/test_task_adjustment.py index 839bbe1..71e3535 100644 --- a/tests/managers/thl/test_task_adjustment.py +++ b/tests/managers/thl/test_task_adjustment.py @@ -1,9 +1,9 @@ import logging +from datetime import UTC, datetime, timedelta, timezone +from decimal import Decimal from random import randint import pytest -from datetime import datetime, timezone, timedelta -from decimal import Decimal from generalresearch.models import Source from generalresearch.models.thl.definitions import ( @@ -31,16 +31,14 @@ def session_complete_with_wallet(session_with_tx_factory, user_with_wallet): @pytest.fixture() def session_fail(user, session_manager, wall_manager): - session = session_manager.create_dummy( - started=datetime.now(timezone.utc), user=user - ) + session = session_manager.create_dummy(started=datetime.now(UTC), user=user) wall1 = wall_manager.create_dummy( session_id=session.id, user_id=user.user_id, source=Source.DYNATA, req_survey_id="72723", req_cpi=Decimal("3.22"), - started=datetime.now(timezone.utc), + started=datetime.now(UTC), ) wall_manager.finish( wall=wall1, @@ -109,7 +107,7 @@ class TestHandleRecons: assert ledger_manager.get_account_balance(commission_account) == 0 # Now, say we get the exact same *adjust to incomplete* msg again. It should do nothing! - adjusted_timestamp = datetime.now(tz=timezone.utc) + adjusted_timestamp = datetime.now(tz=UTC) wall = wall_manager.get_from_uuid(wall_uuid=wall_uuid) with pytest.raises(match=" is already "): wall_manager.adjust_status( diff --git a/tests/managers/thl/test_task_status.py b/tests/managers/thl/test_task_status.py index 55c89c0..468fd5e 100644 --- a/tests/managers/thl/test_task_status.py +++ b/tests/managers/thl/test_task_status.py @@ -1,31 +1,31 @@ -import pytest -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal +import pytest + from generalresearch.managers.thl.session import SessionManager from generalresearch.models import Source from generalresearch.models.thl.definitions import ( Status, - WallAdjustedStatus, StatusCode1, + WallAdjustedStatus, ) from generalresearch.models.thl.product import ( PayoutConfig, - UserWalletConfig, PayoutTransformation, PayoutTransformationPercentArgs, + UserWalletConfig, ) from generalresearch.models.thl.session import Session, WallOut from generalresearch.models.thl.task_status import TaskStatusResponse from generalresearch.models.thl.user import User - -start1 = datetime(2023, 2, 1, tzinfo=timezone.utc) +start1 = datetime(2023, 2, 1, tzinfo=UTC) finish1 = start1 + timedelta(minutes=5) recon1 = start1 + timedelta(days=20) -start2 = datetime(2023, 2, 2, tzinfo=timezone.utc) +start2 = datetime(2023, 2, 2, tzinfo=UTC) finish2 = start2 + timedelta(minutes=5) -start3 = datetime(2023, 2, 3, tzinfo=timezone.utc) +start3 = datetime(2023, 2, 3, tzinfo=UTC) finish3 = start3 + timedelta(minutes=5) diff --git a/tests/managers/thl/test_user_manager/test_base.py b/tests/managers/thl/test_user_manager/test_base.py index 0d7ffef..2704490 100644 --- a/tests/managers/thl/test_user_manager/test_base.py +++ b/tests/managers/thl/test_user_manager/test_base.py @@ -1,5 +1,5 @@ import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from random import randint from uuid import uuid4 @@ -118,7 +118,7 @@ class TestBlockUserManager: ) assert not user.blocked - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) # Adds user to whitelist thl_web_rw.execute_write( """ diff --git a/tests/managers/thl/test_user_streak.py b/tests/managers/thl/test_user_streak.py index 7728f9f..ef25e2b 100644 --- a/tests/managers/thl/test_user_streak.py +++ b/tests/managers/thl/test_user_streak.py @@ -1,17 +1,17 @@ import copy -from datetime import datetime, timezone, timedelta, date +from datetime import UTC, date, datetime, timedelta, timezone from decimal import Decimal from zoneinfo import ZoneInfo import pytest from generalresearch.managers.thl.user_streak import compute_streaks_from_days -from generalresearch.models.thl.definitions import StatusCode1, Status +from generalresearch.models.thl.definitions import Status, StatusCode1 from generalresearch.models.thl.user_streak import ( - UserStreak, - StreakState, - StreakPeriod, StreakFulfillment, + StreakPeriod, + StreakState, + UserStreak, ) @@ -126,7 +126,7 @@ def test_user_streaks_active_broken( user_streak_manager, user, session_manager, broken_active_streak ): # Testing active streak, but broken (not today or yesterday) - start1 = datetime(2025, 2, 12, tzinfo=timezone.utc) + start1 = datetime(2025, 2, 12, tzinfo=UTC) end1 = start1 + timedelta(minutes=1) # abandon counts as inactive @@ -176,7 +176,7 @@ def test_user_streak_complete_active(user_streak_manager, user, session_manager) # They completed yesterday NY time. Today isn't over so streak is pending start1 = datetime.now(tz=ZoneInfo("America/New_York")) - timedelta(days=1) - create_session_complete(session_manager, start1.astimezone(tz=timezone.utc), user) + create_session_complete(session_manager, start1.astimezone(tz=UTC), user) last_complete_day = start1.date() expected_streak = UserStreak( @@ -201,7 +201,7 @@ def test_user_streak_complete_active(user_streak_manager, user, session_manager) # And now they complete today start2 = datetime.now(tz=ZoneInfo("America/New_York")) - create_session_complete(session_manager, start2.astimezone(tz=timezone.utc), user) + create_session_complete(session_manager, start2.astimezone(tz=UTC), user) last_complete_day = start2.date() expected_streak = UserStreak( longest_streak=2, diff --git a/tests/managers/thl/test_userhealth.py b/tests/managers/thl/test_userhealth.py index 1cda8de..e2ea4a8 100644 --- a/tests/managers/thl/test_userhealth.py +++ b/tests/managers/thl/test_userhealth.py @@ -1,4 +1,4 @@ -from datetime import timezone, datetime +from datetime import UTC, datetime, timezone from uuid import uuid4 import faker @@ -12,7 +12,7 @@ from generalresearch.models.thl.ipinfo import GeoIPInformation from generalresearch.models.thl.user_iphistory import ( IPRecord, ) -from generalresearch.models.thl.userhealth import AuditLogLevel, AuditLog +from generalresearch.models.thl.userhealth import AuditLog, AuditLogLevel fake = faker.Faker() @@ -51,7 +51,7 @@ class TestAuditLog: res = audit_log_manager.get_by_id(auditlog_id=audit_log.id) assert isinstance(res, AuditLog) assert res.id == audit_log.id - assert res.created.tzinfo == timezone.utc + assert res.created.tzinfo == UTC def test_filter_by_product( self, @@ -179,7 +179,7 @@ class TestAuditLog: res = audit_log_manager.filter_count( user_ids=[u1.user_id, u2.user_id, u3.user_id], - created_after=datetime.now(tz=timezone.utc), + created_after=datetime.now(tz=UTC), ) assert isinstance(res, int) assert res == 0 diff --git a/tests/managers/thl/test_wall_manager.py b/tests/managers/thl/test_wall_manager.py index ee44e23..46d199f 100644 --- a/tests/managers/thl/test_wall_manager.py +++ b/tests/managers/thl/test_wall_manager.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone, timedelta +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from uuid import uuid4 @@ -10,7 +10,7 @@ from generalresearch.models.thl.session import ( Status, StatusCode1, ) -from test_utils.models.conftest import user, session +from test_utils.models.conftest import session, user class TestWallManager: @@ -88,7 +88,7 @@ class TestWallManager: session_id=session.id, user_id=user.user_id, uuid_id=uuid4().hex, - started=datetime.now(tz=timezone.utc), + started=datetime.now(tz=UTC), source=Source.DYNATA, buyer_id="123", req_survey_id="456", @@ -217,9 +217,9 @@ class TestWallCacheManager: def test_get_wall_events( self, wall_cache_manager, wall_manager, session_manager, user ): - start1 = datetime.now(timezone.utc) - timedelta(hours=3) - start2 = datetime.now(timezone.utc) - timedelta(hours=2) - start3 = datetime.now(timezone.utc) - timedelta(hours=1) + start1 = datetime.now(UTC) - timedelta(hours=3) + start2 = datetime.now(UTC) - timedelta(hours=2) + start3 = datetime.now(UTC) - timedelta(hours=1) session = session_manager.create_dummy(started=start1, user=user) wall1 = wall_manager.create_dummy( diff --git a/tests/models/admin/test_report_request.py b/tests/models/admin/test_report_request.py index a80afbe..4626ab4 100644 --- a/tests/models/admin/test_report_request.py +++ b/tests/models/admin/test_report_request.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone import pandas as pd import pytest @@ -19,19 +19,19 @@ class TestReportRequest: assert rr.report_type == ReportType.POP_SESSION assert rr.start != rr.start_floor, "rr.start != rr.start_floor" - assert rr.start_floor.tzinfo == timezone.utc, "rr.start_floor.tzinfo not utc" + assert rr.start_floor.tzinfo == UTC, "rr.start_floor.tzinfo not utc" rr1 = ReportRequest.model_validate( { "start": datetime( - year=datetime.now(tz=timezone.utc).year, + year=datetime.now(tz=UTC).year, month=1, day=1, hour=0, minute=30, second=25, microsecond=35, - tzinfo=timezone.utc, + tzinfo=UTC, ), "interval": "1h", } @@ -43,14 +43,14 @@ class TestReportRequest: rr2 = ReportRequest.model_validate( { "start": datetime( - year=datetime.now(tz=timezone.utc).year, + year=datetime.now(tz=UTC).year, month=1, day=1, hour=6, minute=30, second=25, microsecond=35, - tzinfo=timezone.utc, + tzinfo=UTC, ), "interval": "1d", } @@ -92,8 +92,8 @@ class TestReportRequest: with pytest.raises(expected_exception=ValidationError): ReportRequest.model_validate( { - "start": datetime(year=1990, month=1, day=1, tzinfo=timezone.utc), - "end": datetime(year=1950, month=1, day=1, tzinfo=timezone.utc), + "start": datetime(year=1990, month=1, day=1, tzinfo=UTC), + "end": datetime(year=1950, month=1, day=1, tzinfo=UTC), } ) @@ -156,8 +156,8 @@ class TestReportRequest: rr = ReportRequest.model_validate( { "interval": "1d", - "start": datetime(year=2000, month=1, day=1, tzinfo=timezone.utc), - "end": datetime(year=2000, month=1, day=10, tzinfo=timezone.utc), + "start": datetime(year=2000, month=1, day=1, tzinfo=UTC), + "end": datetime(year=2000, month=1, day=10, tzinfo=UTC), } ) diff --git a/tests/models/custom_types/test_aware_datetime.py b/tests/models/custom_types/test_aware_datetime.py index 530142e..043fba0 100644 --- a/tests/models/custom_types/test_aware_datetime.py +++ b/tests/models/custom_types/test_aware_datetime.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone import pytest import pytz @@ -27,14 +27,14 @@ class TestAwareDatetimeISO: AwareDatetimeISOModel.model_validate_json(t.model_dump_json()) def test_dt(self): - dt = datetime(2023, 10, 10, 1, 1, 1, tzinfo=timezone.utc) + dt = datetime(2023, 10, 10, 1, 1, 1, tzinfo=UTC) t = AwareDatetimeISOModel(dt=dt, dt_optional=dt) AwareDatetimeISOModel.model_validate_json(t.model_dump_json()) t = AwareDatetimeISOModel(dt=dt, dt_optional=None) AwareDatetimeISOModel.model_validate_json(t.model_dump_json()) - dt = datetime(2023, 10, 10, 1, 1, 1, microsecond=123, tzinfo=timezone.utc) + dt = datetime(2023, 10, 10, 1, 1, 1, microsecond=123, tzinfo=UTC) t = AwareDatetimeISOModel(dt=dt, dt_optional=dt) AwareDatetimeISOModel.model_validate_json(t.model_dump_json()) diff --git a/tests/models/custom_types/test_dsn.py b/tests/models/custom_types/test_dsn.py index 16e1f83..050976e 100644 --- a/tests/models/custom_types/test_dsn.py +++ b/tests/models/custom_types/test_dsn.py @@ -11,9 +11,9 @@ from generalresearch.models.custom_types import DaskDsn, SentryDsn class SettingsModel(BaseModel): - dask: Optional["DaskDsn"] = Field(default=None) - sentry: Optional["SentryDsn"] = Field(default=None) - db: Optional["MySQLDsn"] = Field(default=None) + dask: DaskDsn | None = Field(default=None) + sentry: SentryDsn | None = Field(default=None) + db: MySQLDsn | None = Field(default=None) # --- Pytest themselves --- diff --git a/tests/models/dynata/test_eligbility.py b/tests/models/dynata/test_eligbility.py index 736c971..16cad26 100644 --- a/tests/models/dynata/test_eligbility.py +++ b/tests/models/dynata/test_eligbility.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone class TestEligibility: @@ -40,7 +40,7 @@ class TestEligibility: "project_id": "p1", "status": "OPEN", "project_exclusions": set(), - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), "category_exclusions": set(), "category_ids": set(), "cpi": 1, @@ -172,7 +172,7 @@ class TestEligibility: "project_id": "p1", "status": "OPEN", "project_exclusions": set(), - "created": datetime.now(tz=timezone.utc), + "created": datetime.now(tz=UTC), "category_exclusions": set(), "category_ids": set(), "cpi": 1, diff --git a/tests/models/gr/test_authentication.py b/tests/models/gr/test_authentication.py index 6c84a5d..51595a7 100644 --- a/tests/models/gr/test_authentication.py +++ b/tests/models/gr/test_authentication.py @@ -1,9 +1,9 @@ import binascii import json import os -from datetime import datetime, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timezone from random import randint -from typing import Callable from uuid import uuid4 import pytest @@ -251,7 +251,7 @@ class TestGRToken: def gr_token(self, gr_user): from generalresearch.models.gr.authentication import GRToken - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) token = binascii.hexlify(os.urandom(20)).decode() gr_token = GRToken(key=token, created=now, user_id=gr_user.id) diff --git a/tests/models/gr/test_base.py b/tests/models/gr/test_base.py index a9f01a8..8da28d3 100644 --- a/tests/models/gr/test_base.py +++ b/tests/models/gr/test_base.py @@ -1,6 +1,6 @@ import subprocess +from collections.abc import Callable from pathlib import Path -from typing import Callable import pytest from pydantic import PostgresDsn @@ -10,9 +10,11 @@ from generalresearch.pg_helper import PostgresConfig class TestGRPostgresDjangoCreation: - def test_git(self, git_key_path: Path, gr_repo: Callable[..., Path]): + def test_git(self, gr_repo: Callable[..., Path]): repo_path = gr_repo() + print("test_git.PATH:", repo_path) + try: # Run the git command inside the target directory result = subprocess.run( @@ -36,11 +38,13 @@ class TestGRPostgresDjangoCreation: dsn = django_db_factory("gr") assert isinstance(dsn, PostgresDsn) - # def test_django_tables(self, thl_web_rw: PostgresConfig): - # res = thl_web_rw.execute_sql_query(query=""" - # SELECT COUNT(*) - # FROM information_schema.tables - # WHERE table_schema = 'public'; - # """) - # assert len(res) == 1 - # assert res[0]["count"] == 56 + def test_django_tables(self, gr_db: PostgresConfig): + res = gr_db.execute_sql_query(query=""" + SELECT COUNT(*) + FROM information_schema.tables + WHERE table_schema = 'public'; + """) + print(res) + assert len(res) == 1 + assert res[0]["count"] == 56 + assert True diff --git a/tests/models/gr/test_business.py b/tests/models/gr/test_business.py index 7a84f23..716ec75 100644 --- a/tests/models/gr/test_business.py +++ b/tests/models/gr/test_business.py @@ -1,5 +1,5 @@ import os -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from typing import Optional from uuid import uuid4 @@ -82,15 +82,15 @@ class TestBusinessContact: class TestBusiness: @pytest.fixture - def start(self) -> "datetime": - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + def start(self) -> datetime: + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None def test_init(self, business): @@ -413,15 +413,15 @@ class TestBusiness: class TestBusinessBalance: @pytest.fixture - def start(self) -> "datetime": - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + def start(self) -> datetime: + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None @pytest.mark.skip @@ -1138,7 +1138,7 @@ class TestBusinessBalance: class TestBusinessMethods: @pytest.fixture(scope="function") - def start(self, utc_90days_ago) -> "datetime": + def start(self, utc_90days_ago) -> datetime: s = utc_90days_ago.replace(microsecond=0) return s @@ -1149,7 +1149,7 @@ class TestBusinessMethods: @pytest.fixture(scope="function") def duration( self, - ) -> Optional["timedelta"]: + ) -> timedelta | None: return None def test_cache_key(self, business, gr_redis): @@ -1212,7 +1212,7 @@ class TestBusinessMethods: # We're going to pull only a specific year, but make sure that # it's being assigned to the field regardless - year = datetime.now(tz=timezone.utc).year + year = datetime.now(tz=UTC).year res = Business.from_redis( uuid=business.uuid, fields=[f"pop_financial:{year}"], diff --git a/tests/models/morning/test.py b/tests/models/morning/test.py index bedf9c2..222cb93 100644 --- a/tests/models/morning/test.py +++ b/tests/models/morning/test.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from generalresearch.models.morning.question import MorningQuestion @@ -163,8 +163,8 @@ bid = { # what gets run in MorningAPI._format_bid bid["language_isos"] = ("eng",) bid["country_iso"] = "us" -bid["end_date"] = datetime(2024, 7, 19, 9, 1, 13, 520243, tzinfo=timezone.utc) -bid["published_at"] = datetime(2024, 6, 19, 9, 1, 13, 520243, tzinfo=timezone.utc) +bid["end_date"] = datetime(2024, 7, 19, 9, 1, 13, 520243, tzinfo=UTC) +bid["published_at"] = datetime(2024, 6, 19, 9, 1, 13, 520243, tzinfo=UTC) bid.update(bid["statistics"]) bid["qualified_conversion"] /= 100 bid["system_conversion"] /= 100 diff --git a/tests/models/prodege/test_survey_participation.py b/tests/models/prodege/test_survey_participation.py index 68d7838..3b35d0c 100644 --- a/tests/models/prodege/test_survey_participation.py +++ b/tests/models/prodege/test_survey_participation.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone class TestProdegeParticipation: @@ -10,7 +10,7 @@ class TestProdegeParticipation: ProdegeUserPastParticipation, ) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) pp = ProdegePastParticipation.from_api( { "participation_project_ids": [152677146, 152803285], @@ -89,7 +89,7 @@ class TestProdegeParticipation: ProdegeUserPastParticipation, ) - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) pp = ProdegePastParticipation.from_api( { "participation_project_ids": [152677146, 152803285], diff --git a/tests/models/spectrum/test_question.py b/tests/models/spectrum/test_question.py index ba118d7..4f92961 100644 --- a/tests/models/spectrum/test_question.py +++ b/tests/models/spectrum/test_question.py @@ -1,17 +1,17 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from generalresearch.models import Source from generalresearch.models.spectrum.question import ( - SpectrumQuestionOption, SpectrumQuestion, - SpectrumQuestionType, SpectrumQuestionClass, + SpectrumQuestionOption, + SpectrumQuestionType, ) from generalresearch.models.thl.profiling.upk_question import ( UpkQuestion, + UpkQuestionChoice, UpkQuestionSelectorMC, UpkQuestionType, - UpkQuestionChoice, ) @@ -43,7 +43,7 @@ class TestSpectrumQuestion: tags=None, options=None, class_num=SpectrumQuestionClass.CORE, - created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc), + created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC), is_live=True, source=Source.SPECTRUM, category_id=None, @@ -85,7 +85,7 @@ class TestSpectrumQuestion: SpectrumQuestionOption(id="112", text="Female", order=1), ], class_num=SpectrumQuestionClass.CORE, - created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc), + created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC), is_live=True, source=Source.SPECTRUM, category_id=None, @@ -160,7 +160,7 @@ class TestSpectrumQuestion: SpectrumQuestionOption(id="999", text="None of the above", order=3), ], class_num=SpectrumQuestionClass.EXTENDED, - created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=timezone.utc), + created=datetime(2017, 8, 16, 7, 52, 7, 688000, tzinfo=UTC), is_live=True, source=Source.SPECTRUM, category_id=None, diff --git a/tests/models/spectrum/test_survey.py b/tests/models/spectrum/test_survey.py index b612a63..5e095a3 100644 --- a/tests/models/spectrum/test_survey.py +++ b/tests/models/spectrum/test_survey.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from decimal import Decimal @@ -140,11 +140,11 @@ class TestSpectrumSurvey: "survey_id": 29333264, "survey_name": "Exciting New Survey #29333264", "survey_status": 22, - "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc), + "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC), "category": "Exciting New", "category_code": 232, - "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc), - "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc), + "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC), + "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC), "soft_launch": False, "click_balancing": 0, "price_type": 1, @@ -212,7 +212,7 @@ class TestSpectrumSurvey: survey_id="29333264", survey_name="Exciting New Survey #29333264", status=SpectrumStatus.LIVE, - field_end_date=datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc), + field_end_date=datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC), category_code="232", calculation_type=TaskCalculationType.COMPLETES, requires_pii=False, @@ -240,8 +240,8 @@ class TestSpectrumSurvey: values=["18-64"], ) }, - created_api=datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc), - modified_api=datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc), + created_api=datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC), + modified_api=datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC), updated=None, ) assert expected_survey.model_dump_json() == s.model_dump_json() @@ -255,11 +255,11 @@ class TestSpectrumSurvey: "survey_id": 29333264, "survey_name": "#29333264", "survey_status": 22, - "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc), + "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC), "category": "Exciting New", "category_code": 232, - "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc), - "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc), + "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC), + "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC), "soft_launch": False, "click_balancing": 0, "price_type": 1, @@ -318,11 +318,11 @@ class TestSpectrumSurvey: "survey_id": 29333264, "survey_name": "#29333264", "survey_status": 22, - "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc), + "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC), "category": "Exciting New", "category_code": 232, - "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc), - "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc), + "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC), + "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC), "soft_launch": False, "click_balancing": 0, "price_type": 1, diff --git a/tests/models/spectrum/test_survey_manager.py b/tests/models/spectrum/test_survey_manager.py index 582093c..11970bf 100644 --- a/tests/models/spectrum/test_survey_manager.py +++ b/tests/models/spectrum/test_survey_manager.py @@ -1,22 +1,21 @@ import copy import logging -from datetime import timezone, datetime +from datetime import UTC, datetime, timezone from decimal import Decimal from pymysql import IntegrityError - logger = logging.getLogger() example_survey_api_response = { "survey_id": 29333264, "survey_name": "#29333264", "survey_status": 22, - "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=timezone.utc), + "field_end_date": datetime(2024, 5, 23, 18, 18, 31, tzinfo=UTC), "category": "Exciting New", "category_code": 232, - "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=timezone.utc), - "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=timezone.utc), + "crtd_on": datetime(2024, 5, 20, 17, 48, 13, tzinfo=UTC), + "mod_on": datetime(2024, 5, 20, 18, 18, 31, tzinfo=UTC), "soft_launch": False, "click_balancing": 0, "price_type": 1, @@ -66,7 +65,7 @@ class TestSpectrumSurvey: assert settings.debug, "CRITICAL: Do not run this on production." - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) spectrum_rw.execute_sql_query( query=f""" DELETE FROM `{spectrum_rw.db}`.spectrum_survey @@ -93,7 +92,7 @@ class TestSpectrumSurvey: assert settings.debug, "CRITICAL: Do not run this on production." - now = datetime.now(tz=timezone.utc) + now = datetime.now(tz=UTC) spectrum_rw.execute_sql_query( query=f""" DELETE FROM `{spectrum_rw.db}`.spectrum_survey diff --git a/tests/models/test_finance.py b/tests/models/test_finance.py index bd548b3..6dcd441 100644 --- a/tests/models/test_finance.py +++ b/tests/models/test_finance.py @@ -1,7 +1,7 @@ -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from itertools import product as iter_product from random import randint -from typing import Callable from uuid import uuid4 import pandas as pd @@ -683,7 +683,7 @@ class TestProductFinanceData: rand_item_time = fake.date_time_between( start_date=item.start, end_date=item.finish, - tzinfo=timezone.utc, + tzinfo=UTC, ) session_with_tx_factory(started=rand_item_time, user=u) @@ -773,7 +773,7 @@ class TestPOPFinancialData: rand_item_time = fake.date_time_between( start_date=item.start, end_date=item.finish, - tzinfo=timezone.utc, + tzinfo=UTC, ) session_with_tx_factory(started=rand_item_time, user=u) @@ -870,7 +870,7 @@ class TestBusinessBalanceData: item_time = fake.date_time_between( start_date=item.start, end_date=item.finish, - tzinfo=timezone.utc, + tzinfo=UTC, ) session_with_tx_factory(started=item_time, user=u) item.initial_load(overwrite=True) diff --git a/tests/models/thl/test_adjustments.py b/tests/models/thl/test_adjustments.py index 27091bb..c2c035d 100644 --- a/tests/models/thl/test_adjustments.py +++ b/tests/models/thl/test_adjustments.py @@ -1,6 +1,6 @@ -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import Callable import pytest @@ -16,14 +16,14 @@ from generalresearch.models.thl.session import ( ) from generalresearch.models.thl.user import User -started1 = datetime(2023, 1, 1, tzinfo=timezone.utc) -started2 = datetime(2023, 1, 1, 0, 10, 0, tzinfo=timezone.utc) +started1 = datetime(2023, 1, 1, tzinfo=UTC) +started2 = datetime(2023, 1, 1, 0, 10, 0, tzinfo=UTC) finished1 = started1 + timedelta(minutes=10) finished2 = started2 + timedelta(minutes=10) -adj_ts = datetime(2023, 2, 2, tzinfo=timezone.utc) -adj_ts2 = datetime(2023, 2, 3, tzinfo=timezone.utc) -adj_ts3 = datetime(2023, 2, 4, tzinfo=timezone.utc) +adj_ts = datetime(2023, 2, 2, tzinfo=UTC) +adj_ts2 = datetime(2023, 2, 3, tzinfo=UTC) +adj_ts3 = datetime(2023, 2, 4, tzinfo=UTC) class TestProductAdjustments: diff --git a/tests/models/thl/test_contest/test_contest.py b/tests/models/thl/test_contest/test_contest.py index 0fbd4cc..acb501c 100644 --- a/tests/models/thl/test_contest/test_contest.py +++ b/tests/models/thl/test_contest/test_contest.py @@ -1,4 +1,4 @@ -from typing import Callable +from collections.abc import Callable import pytest diff --git a/tests/models/thl/test_contest/test_leaderboard_contest.py b/tests/models/thl/test_contest/test_leaderboard_contest.py index 8b714ee..3efcf2f 100644 --- a/tests/models/thl/test_contest/test_leaderboard_contest.py +++ b/tests/models/thl/test_contest/test_leaderboard_contest.py @@ -1,4 +1,4 @@ -from datetime import timezone +from datetime import UTC, timezone from uuid import uuid4 import pytest @@ -26,7 +26,7 @@ class TestLeaderboardContest(TestContest): @pytest.fixture def leaderboard_contest( self, product: Product, thl_redis, user_manager - ) -> "LeaderboardContest": + ) -> LeaderboardContest: board_key = f"leaderboard:{product.uuid}:us:weekly:2025-05-26:complete_count" c = LeaderboardContest( @@ -91,7 +91,7 @@ class TestLeaderboardContest(TestContest): country_iso=model.country_iso, freq=model.freq, product_id=leaderboard_contest.product_id, - within_time=model.period_start_local.astimezone(tz=timezone.utc), + within_time=model.period_start_local.astimezone(tz=UTC), ) lbm.hit_complete_count(product_user_id=user_1.product_user_id) diff --git a/tests/models/thl/test_ledger.py b/tests/models/thl/test_ledger.py index 257de3c..5edcc9d 100644 --- a/tests/models/thl/test_ledger.py +++ b/tests/models/thl/test_ledger.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import UTC, datetime, timezone from uuid import uuid4 import pytest @@ -21,7 +21,7 @@ class TestLedgerTransaction: assert [] == t.entries assert {} == t.metadata t = LedgerTransaction( - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), metadata={"a": "b", "user": "1234"}, ext_description="foo", ) diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index 39469dc..78bc10a 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -2,9 +2,9 @@ from __future__ import annotations import os import shutil -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal -from typing import Callable from uuid import uuid4 import pytest @@ -586,7 +586,7 @@ class TestProductFinancials: @pytest.fixture def start(self) -> datetime: - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: @@ -769,7 +769,7 @@ class TestProductBalance: @pytest.fixture def start(self) -> datetime: - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: @@ -877,7 +877,7 @@ class TestProductBalance: product=product, amount=USDCent(71), ext_ref_id=uuid4().hex, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), skip_wallet_balance_check=True, skip_one_per_day_check=True, ) @@ -892,7 +892,7 @@ class TestProductPOPFinancial: @pytest.fixture def start(self) -> datetime: - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: @@ -965,7 +965,7 @@ class TestProductCache: @pytest.fixture def start(self) -> datetime: - return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) + return datetime(year=2018, month=3, day=14, hour=0, tzinfo=UTC) @pytest.fixture def offset(self) -> str: diff --git a/tests/models/thl/test_user.py b/tests/models/thl/test_user.py index 943ae8e..a4f331a 100644 --- a/tests/models/thl/test_user.py +++ b/tests/models/thl/test_user.py @@ -1,5 +1,5 @@ import json -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from random import choice as rand_choice from random import randint @@ -383,7 +383,7 @@ class TestUserCreated: from generalresearch.models.thl.user import User user = User(user_id=self.user_id) - dt = datetime.now(tz=timezone.utc) + dt = datetime.now(tz=UTC) user.created = dt assert user.created == dt @@ -419,7 +419,7 @@ class TestUserCreated: def test_not_in_future(self): from generalresearch.models.thl.user import User - the_future = datetime.now(tz=timezone.utc) + timedelta(minutes=1) + the_future = datetime.now(tz=UTC) + timedelta(minutes=1) with pytest.raises(ValueError) as cm: User(user_id=self.user_id, created=the_future) assert "1 validation error for User" in str(cm.value) @@ -428,9 +428,9 @@ class TestUserCreated: def test_after_anno_domini(self): from generalresearch.models.thl.user import User - before_ad = datetime( - year=2015, month=1, day=1, tzinfo=timezone.utc - ) + timedelta(minutes=1) + before_ad = datetime(year=2015, month=1, day=1, tzinfo=UTC) + timedelta( + minutes=1 + ) with pytest.raises(ValueError) as cm: User(user_id=self.user_id, created=before_ad) assert "1 validation error for User" in str(cm.value) @@ -444,7 +444,7 @@ class TestUserLastSeen: from generalresearch.models.thl.user import User user = User(user_id=self.user_id) - dt = datetime.now(tz=timezone.utc) + dt = datetime.now(tz=UTC) user.last_seen = dt assert user.last_seen == dt @@ -480,7 +480,7 @@ class TestUserLastSeen: def test_not_in_future(self): from generalresearch.models.thl.user import User - the_future = datetime.now(tz=timezone.utc) + timedelta(minutes=1) + the_future = datetime.now(tz=UTC) + timedelta(minutes=1) with pytest.raises(ValueError) as cm: User(user_id=self.user_id, last_seen=the_future) assert "1 validation error for User" in str(cm.value) @@ -489,9 +489,9 @@ class TestUserLastSeen: def test_after_anno_domini(self): from generalresearch.models.thl.user import User - before_ad = datetime( - year=2015, month=1, day=1, tzinfo=timezone.utc - ) + timedelta(minutes=1) + before_ad = datetime(year=2015, month=1, day=1, tzinfo=UTC) + timedelta( + minutes=1 + ) with pytest.raises(ValueError) as cm: User(user_id=self.user_id, last_seen=before_ad) assert "1 validation error for User" in str(cm.value) @@ -549,8 +549,8 @@ class TestUserTiming: def test_valid(self): from generalresearch.models.thl.user import User - created = datetime.now(tz=timezone.utc) - timedelta(minutes=60) - last_seen = datetime.now(tz=timezone.utc) - timedelta(minutes=59) + created = datetime.now(tz=UTC) - timedelta(minutes=60) + last_seen = datetime.now(tz=UTC) - timedelta(minutes=59) user = User(user_id=self.user_id, created=created, last_seen=last_seen) assert user.created == created @@ -559,8 +559,8 @@ class TestUserTiming: def test_created_first(self): from generalresearch.models.thl.user import User - created = datetime.now(tz=timezone.utc) - timedelta(minutes=60) - last_seen = datetime.now(tz=timezone.utc) - timedelta(minutes=59) + created = datetime.now(tz=UTC) - timedelta(minutes=60) + last_seen = datetime.now(tz=UTC) - timedelta(minutes=59) with pytest.raises(ValueError) as cm: User(user_id=self.user_id, created=last_seen, last_seen=created) @@ -602,7 +602,7 @@ class TestUserSerialization: user = User( product_id=product_id, product_user_id=product_user_id, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), blocked=False, ) @@ -623,7 +623,7 @@ class TestUserSerialization: user = User( product_id=product_id, product_user_id=product_user_id, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), blocked=False, ) @@ -633,7 +633,7 @@ class TestUserSerialization: assert not d.get("blocked") assert d.get("product") is None - assert d.get("created").tzinfo == timezone.utc + assert d.get("created").tzinfo == UTC def test_from_json(self): from generalresearch.models.thl.user import User @@ -644,14 +644,14 @@ class TestUserSerialization: user = User( product_id=product_id, product_user_id=product_user_id, - created=datetime.now(tz=timezone.utc), + created=datetime.now(tz=UTC), blocked=False, ) u = User.model_validate_json(user.to_json()) assert u.product_id == product_id assert u.product is None - assert u.created.tzinfo == timezone.utc + assert u.created.tzinfo == UTC class TestUserMethods: diff --git a/tests/models/thl/test_user_iphistory.py b/tests/models/thl/test_user_iphistory.py index 596849c..0f050b0 100644 --- a/tests/models/thl/test_user_iphistory.py +++ b/tests/models/thl/test_user_iphistory.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from generalresearch.models.thl.user_iphistory import ( UserIPHistory, @@ -8,7 +8,7 @@ from generalresearch.models.thl.user_iphistory import ( def test_collapse_ip_records(): # This does not exist in a db, so we do not need fixtures/ real user ids, whatever - now = datetime.now(tz=timezone.utc) - timedelta(days=1) + now = datetime.now(tz=UTC) - timedelta(days=1) # Gets stored most recent first. This is reversed, but the validator will order it records = [ UserIPRecord(ip="1.2.3.5", created=now + timedelta(minutes=1)), diff --git a/tests/models/thl/test_wall.py b/tests/models/thl/test_wall.py index 8398c81..9e9483b 100644 --- a/tests/models/thl/test_wall.py +++ b/tests/models/thl/test_wall.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from uuid import uuid4 @@ -27,8 +27,8 @@ class TestWall: ext_status_code_1="1.0", status=Status.FAIL, status_code_1=StatusCode1.BUYER_FAIL, - started=datetime(2023, 1, 1, 0, 0, 1, tzinfo=timezone.utc), - finished=datetime(2023, 1, 1, 0, 10, 1, tzinfo=timezone.utc), + started=datetime(2023, 1, 1, 0, 0, 1, tzinfo=UTC), + finished=datetime(2023, 1, 1, 0, 10, 1, tzinfo=UTC), ) s = w.to_json() w2 = Wall.from_json(s) @@ -45,8 +45,8 @@ class TestWall: survey_id="yyy", status=Status.FAIL, status_code_1=StatusCode1.BUYER_FAIL, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) Wall( user_id=1, @@ -58,8 +58,8 @@ class TestWall: status=Status.FAIL, status_code_1=StatusCode1.MARKETPLACE_FAIL, status_code_2=WallStatusCode2.COMPLETE_TOO_FAST, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) with pytest.raises(expected_exception=ValidationError) as e: Wall( @@ -71,8 +71,8 @@ class TestWall: survey_id="yyy", status=Status.FAIL, status_code_1=StatusCode1.GRS_ABANDON, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) assert "If status is f, status_code_1 should be in" in str(e.value) @@ -87,8 +87,8 @@ class TestWall: status=Status.FAIL, status_code_1=StatusCode1.GRS_ABANDON, status_code_2=WallStatusCode2.COMPLETE_TOO_FAST, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) assert "If status is f, status_code_1 should be in" in str(e.value) @@ -104,8 +104,8 @@ class TestWall: status=Status.FAIL, status_code_1=StatusCode1.MARKETPLACE_FAIL, status_code_2=WallStatusCode2.COMPLETE_TOO_FAST, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) Wall( user_id=1, @@ -117,8 +117,8 @@ class TestWall: status=Status.FAIL, status_code_1=StatusCode1.BUYER_FAIL, status_code_2=None, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) Wall( user_id=1, @@ -130,8 +130,8 @@ class TestWall: status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, status_code_2=None, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) with pytest.raises(expected_exception=ValidationError) as e: @@ -145,8 +145,8 @@ class TestWall: status=Status.FAIL, status_code_1=StatusCode1.BUYER_FAIL, status_code_2=WallStatusCode2.COMPLETE_TOO_FAST, - started=datetime.now(timezone.utc), - finished=datetime.now(timezone.utc) + timedelta(seconds=1), + started=datetime.now(UTC), + finished=datetime.now(UTC) + timedelta(seconds=1), ) assert "If status_code_1 is 1, status_code_2 should be in" in str(e.value) diff --git a/tests/models/thl/test_wall_session.py b/tests/models/thl/test_wall_session.py index 1208c56..10f3cba 100644 --- a/tests/models/thl/test_wall_session.py +++ b/tests/models/thl/test_wall_session.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal import pytest @@ -12,7 +12,7 @@ from generalresearch.models.thl.user import User class TestWallSession: def test_session_with_no_wall_events(self): - started = datetime(2023, 1, 1, tzinfo=timezone.utc) + started = datetime(2023, 1, 1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) assert s.status is None assert s.status_code_1 is None @@ -24,7 +24,7 @@ class TestWallSession: # assert s.status_code_1 == StatusCode1.SESSION_START_FAIL def test_session_timeout_with_only_grs(self): - started = datetime(2023, 1, 1, tzinfo=timezone.utc) + started = datetime(2023, 1, 1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( user_id=1, @@ -53,7 +53,7 @@ class TestWallSession: # assert s.status_code_1 == StatusCode1.GRS_FAIL def test_session_with_only_grs_complete(self): - started = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc) + started = datetime(year=2023, month=1, day=1, tzinfo=UTC) # A Session is started s = Session(user=User(user_id=1), started=started) @@ -98,7 +98,7 @@ class TestWallSession: # assert s.status_code_1 is None def test_session_with_only_non_grs_fail(self): - started = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc) + started = datetime(year=2023, month=1, day=1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( @@ -119,7 +119,7 @@ class TestWallSession: assert s.payout is None def test_session_with_only_non_grs_timeout(self): - started = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc) + started = datetime(year=2023, month=1, day=1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( @@ -139,7 +139,7 @@ class TestWallSession: assert s.payout is None def test_session_with_grs_and_external(self): - started = datetime(year=2023, month=1, day=1, tzinfo=timezone.utc) + started = datetime(year=2023, month=1, day=1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( @@ -168,7 +168,7 @@ class TestWallSession: s.append_wall_event(w) w.finish( status=Status.ABANDON, - finished=datetime.now(tz=timezone.utc) + timedelta(minutes=10), + finished=datetime.now(tz=UTC) + timedelta(minutes=10), status_code_1=StatusCode1.BUYER_ABANDON, ) status, status_code_1 = s.determine_session_status() @@ -206,7 +206,7 @@ class TestWallSession: assert s.payout is None def test_session_marketplace_fail(self): - started = datetime(2023, 1, 1, tzinfo=timezone.utc) + started = datetime(2023, 1, 1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( @@ -229,7 +229,7 @@ class TestWallSession: assert StatusCode1.SESSION_CONTINUE_QUALITY_FAIL == s.status_code_1 def test_session_unknown(self): - started = datetime(2023, 1, 1, tzinfo=timezone.utc) + started = datetime(2023, 1, 1, tzinfo=UTC) s = Session(user=User(user_id=1), started=started) w = Wall( diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 3b3ddd0..ed5a7ae 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -1,6 +1,6 @@ import socket import subprocess -from typing import Callable +from collections.abc import Callable from pydantic import PostgresDsn @@ -12,7 +12,7 @@ def is_port_open(host: InternalHostname, port: int = 5432, timeout: int = 3): try: with socket.create_connection((host, port), timeout=timeout): return True - except (socket.timeout, ConnectionRefusedError, OSError): + except (TimeoutError, ConnectionRefusedError, OSError): return False -- cgit v1.2.3 From e2c5de703be45746bacaea4136f24440ff5a291c Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Mon, 24 Aug 2026 12:35:35 -0700 Subject: Ruff std replacements --- generalresearch/__init__.py | 2 +- generalresearch/config.py | 2 +- generalresearch/grliq/managers/event_plotter.py | 4 +- generalresearch/grliq/managers/forensic_data.py | 4 +- generalresearch/grliq/managers/forensic_events.py | 76 ++++++++++------------ generalresearch/grliq/managers/forensic_results.py | 7 +- generalresearch/grliq/models/custom_types.py | 3 +- generalresearch/grliq/models/forensic_data.py | 1 - generalresearch/grliq/models/forensic_summary.py | 3 +- generalresearch/grliq/utils.py | 2 +- generalresearch/grpc.py | 2 +- generalresearch/incite/base.py | 10 +-- generalresearch/incite/collections/__init__.py | 18 ++--- generalresearch/incite/defaults.py | 2 +- .../incite/mergers/foundations/__init__.py | 2 +- .../incite/mergers/foundations/enriched_session.py | 10 +-- .../mergers/foundations/enriched_task_adjust.py | 5 +- .../incite/mergers/foundations/enriched_wall.py | 8 +-- generalresearch/incite/mergers/ym_survey_wall.py | 10 ++- generalresearch/incite/mergers/ym_wall_summary.py | 3 +- generalresearch/incite/schemas/thl_web.py | 2 +- generalresearch/locales/setup_json.py | 1 + generalresearch/locales/timezone.py | 1 - generalresearch/managers/cint/survey.py | 2 +- generalresearch/managers/criteria.py | 4 +- generalresearch/managers/dynata/survey.py | 2 +- generalresearch/managers/events.py | 4 +- generalresearch/managers/gr/authentication.py | 5 +- generalresearch/managers/gr/team.py | 2 +- generalresearch/managers/innovate/survey.py | 2 +- generalresearch/managers/leaderboard/__init__.py | 3 +- generalresearch/managers/leaderboard/manager.py | 2 +- generalresearch/managers/morning/survey.py | 2 +- generalresearch/managers/network/label.py | 9 ++- generalresearch/managers/network/tool_run.py | 22 +++---- generalresearch/managers/precision/survey.py | 10 +-- generalresearch/managers/prodege/survey.py | 2 +- generalresearch/managers/repdata/survey.py | 4 +- generalresearch/managers/sago/survey.py | 2 +- generalresearch/managers/spectrum/survey.py | 2 +- generalresearch/managers/survey.py | 2 - generalresearch/managers/thl/buyer.py | 2 +- generalresearch/managers/thl/cashout_method.py | 2 +- generalresearch/managers/thl/contest_manager.py | 3 +- .../managers/thl/ledger_manager/conditions.py | 4 +- .../managers/thl/ledger_manager/exceptions.py | 5 -- .../managers/thl/ledger_manager/ledger.py | 15 ++--- .../managers/thl/ledger_manager/thl_ledger.py | 22 +++---- generalresearch/managers/thl/product.py | 10 +-- generalresearch/managers/thl/profiling/question.py | 2 +- generalresearch/managers/thl/profiling/uqa.py | 2 +- generalresearch/managers/thl/profiling/user_upk.py | 2 +- generalresearch/managers/thl/session.py | 11 +--- generalresearch/managers/thl/survey.py | 25 ++++--- generalresearch/managers/thl/task_adjustment.py | 7 +- generalresearch/managers/thl/user_compensate.py | 2 +- .../thl/user_manager/mysql_user_manager.py | 2 +- .../managers/thl/user_manager/user_manager.py | 1 - generalresearch/managers/thl/userhealth.py | 2 +- generalresearch/managers/thl/wall.py | 2 +- generalresearch/managers/thl/wallet/tango.py | 2 +- generalresearch/models/admin/request.py | 2 +- generalresearch/models/custom_types.py | 3 +- generalresearch/models/dynata/survey.py | 4 +- generalresearch/models/gr/authentication.py | 5 +- generalresearch/models/gr/business.py | 6 +- generalresearch/models/gr/team.py | 7 +- generalresearch/models/legacy/questions.py | 3 - generalresearch/models/network/mtr/execute.py | 2 +- generalresearch/models/network/nmap/parser.py | 2 +- generalresearch/models/network/rdns/execute.py | 2 +- generalresearch/models/spectrum/survey.py | 2 +- generalresearch/models/string_utils.py | 1 - generalresearch/models/thl/contest/contest.py | 4 +- .../models/thl/contest/contest_entry.py | 2 +- generalresearch/models/thl/contest/examples.py | 7 -- generalresearch/models/thl/contest/io.py | 2 +- generalresearch/models/thl/contest/leaderboard.py | 2 +- generalresearch/models/thl/contest/milestone.py | 3 +- generalresearch/models/thl/contest/raffle.py | 2 +- generalresearch/models/thl/finance.py | 5 +- generalresearch/models/thl/ipinfo.py | 2 +- generalresearch/models/thl/ledger_example.py | 2 +- generalresearch/models/thl/offerwall/cache.py | 2 +- generalresearch/models/thl/payout_format.py | 12 ++-- generalresearch/models/thl/product.py | 2 - .../models/thl/profiling/marketplace.py | 2 +- .../models/thl/profiling/upk_question.py | 2 +- .../models/thl/profiling/upk_question_answer.py | 2 +- .../models/thl/profiling/user_question_answer.py | 2 +- generalresearch/models/thl/session.py | 2 +- generalresearch/models/thl/survey/__init__.py | 3 - generalresearch/models/thl/survey/buyer.py | 2 +- generalresearch/models/thl/survey/model.py | 2 +- generalresearch/models/thl/survey/penalty.py | 2 +- generalresearch/models/thl/task_adjustment.py | 2 +- generalresearch/models/thl/user.py | 2 +- generalresearch/models/thl/user_iphistory.py | 3 +- generalresearch/models/thl/wallet/payout.py | 2 +- generalresearch/pg_helper.py | 12 ++-- generalresearch/schemas/survey_stats.py | 2 +- generalresearch/sql_helper.py | 19 ++---- generalresearch/thl_django/apps.py | 14 ++-- generalresearch/thl_django/fields.py | 3 +- .../thl_django/migrations/0001_initial.py | 3 +- ..._live_alter_surveycategory_strength_and_more.py | 2 +- ...rveystat_surveystat_live_survey_idx_and_more.py | 2 +- ...ssion_thl_session_status_d578b7_idx_and_more.py | 2 +- ...p_portscanport_iplabel_mtr_portscan_and_more.py | 5 +- generalresearch/thl_django/network/models.py | 5 +- generalresearch/utils/enum.py | 4 +- generalresearch/wall_status_codes/lucid.py | 2 +- generalresearch/wall_status_codes/morning.py | 2 +- generalresearch/wall_status_codes/pollfish.py | 2 +- generalresearch/wall_status_codes/precision.py | 2 +- generalresearch/wall_status_codes/repdata.py | 2 +- test_utils/conftest.py | 3 +- test_utils/grliq/conftest.py | 2 +- test_utils/incite/conftest.py | 2 +- test_utils/managers/conftest.py | 2 - test_utils/models/conftest.py | 9 +-- test_utils/models/contest/conftest.py | 2 +- test_utils/models/network/conftest.py | 2 +- test_utils/models/thl/conftest.py | 2 +- test_utils/spectrum/conftest.py | 2 +- .../incite/collections/test_df_collection_base.py | 3 +- .../collections/test_df_collection_item_base.py | 2 +- .../collections/test_df_collection_item_thl_web.py | 11 +--- .../test_df_collection_thl_marketplaces.py | 7 +- .../collections/test_df_collection_thl_web.py | 2 - .../mergers/foundations/test_enriched_session.py | 7 +- .../foundations/test_enriched_task_adjust.py | 7 -- .../mergers/foundations/test_enriched_wall.py | 23 +------ .../mergers/foundations/test_user_id_product.py | 15 +---- tests/incite/mergers/test_merge_collection.py | 3 +- tests/incite/mergers/test_merge_collection_item.py | 9 +-- tests/incite/mergers/test_pop_ledger.py | 8 +-- tests/incite/mergers/test_ym_survey_merge.py | 17 +---- tests/incite/test_collection_base.py | 3 +- tests/incite/test_collection_base_item.py | 2 +- tests/incite/test_grl_flow.py | 11 ++-- tests/incite/test_interval_idx.py | 3 +- tests/managers/gr/test_authentication.py | 3 - tests/managers/gr/test_business.py | 4 +- tests/managers/gr/test_team.py | 2 - tests/managers/leaderboard.py | 2 +- tests/managers/network/test_label.py | 2 +- tests/managers/test_events.py | 15 ++--- tests/managers/test_userpid.py | 4 +- .../managers/thl/test_contest/test_leaderboard.py | 8 +-- tests/managers/thl/test_contest/test_milestone.py | 14 +--- tests/managers/thl/test_contest/test_raffle.py | 14 +--- tests/managers/thl/test_harmonized_uqa.py | 2 +- tests/managers/thl/test_ipinfo.py | 2 +- tests/managers/thl/test_ledger/test_lm_accounts.py | 6 -- tests/managers/thl/test_ledger/test_lm_tx_locks.py | 23 ++----- .../thl/test_ledger/test_thl_lm_accounts.py | 32 ++++----- .../thl/test_ledger/test_thl_lm_bp_payout.py | 10 +-- tests/managers/thl/test_ledger/test_thl_lm_tx.py | 64 +++--------------- .../test_ledger/test_thl_lm_tx__user_payouts.py | 3 +- tests/managers/thl/test_ledger/test_thl_pem.py | 7 +- tests/managers/thl/test_ledger/test_user_txs.py | 3 +- tests/managers/thl/test_ledger/test_wallet.py | 2 +- tests/managers/thl/test_product.py | 9 ++- tests/managers/thl/test_product_prod.py | 2 - tests/managers/thl/test_profiling/test_user_upk.py | 2 +- tests/managers/thl/test_session_manager.py | 7 +- tests/managers/thl/test_survey.py | 2 +- tests/managers/thl/test_survey_penalty.py | 1 - tests/managers/thl/test_task_adjustment.py | 2 +- tests/managers/thl/test_task_status.py | 2 +- tests/managers/thl/test_user_manager/test_base.py | 4 +- tests/managers/thl/test_user_manager/test_mysql.py | 1 - .../thl/test_user_manager/test_user_fetch.py | 1 - .../thl/test_user_manager/test_user_metadata.py | 1 - tests/managers/thl/test_user_streak.py | 2 +- tests/managers/thl/test_userhealth.py | 2 +- tests/managers/thl/test_wall_manager.py | 13 ++-- tests/models/admin/test_report_request.py | 2 +- tests/models/custom_types/test_aware_datetime.py | 2 +- tests/models/custom_types/test_dsn.py | 1 - tests/models/dynata/test_eligbility.py | 2 +- tests/models/gr/test_authentication.py | 2 +- tests/models/gr/test_business.py | 3 +- tests/models/innovate/test_question.py | 6 +- .../models/legacy/test_user_question_answer_in.py | 4 +- tests/models/morning/test.py | 2 +- tests/models/network/test_mtr.py | 4 +- tests/models/network/test_nmap_parser.py | 1 + tests/models/prodege/test_survey_participation.py | 2 +- tests/models/spectrum/test_question.py | 2 +- tests/models/spectrum/test_survey.py | 2 +- tests/models/spectrum/test_survey_manager.py | 2 +- tests/models/test_finance.py | 2 +- tests/models/thl/question/test_question_info.py | 2 +- tests/models/thl/test_adjustments.py | 30 ++------- .../thl/test_contest/test_leaderboard_contest.py | 2 +- tests/models/thl/test_ledger.py | 2 +- tests/models/thl/test_payout.py | 5 +- tests/models/thl/test_product.py | 2 +- tests/models/thl/test_upkquestion.py | 12 +--- tests/models/thl/test_user.py | 6 +- tests/models/thl/test_user_iphistory.py | 2 +- tests/models/thl/test_user_streak.py | 2 +- tests/models/thl/test_wall.py | 2 +- tests/models/thl/test_wall_session.py | 2 +- tests/sql_helper.py | 4 +- 207 files changed, 398 insertions(+), 725 deletions(-) (limited to 'test_utils/models/network') diff --git a/generalresearch/__init__.py b/generalresearch/__init__.py index 604b7e2..3b2ec3d 100644 --- a/generalresearch/__init__.py +++ b/generalresearch/__init__.py @@ -129,7 +129,7 @@ def synchronized(wrapped): if lock is None: lock = threading.RLock() - setattr(context, "_synchronized_lock", lock) + context._synchronized_lock = lock return lock diff --git a/generalresearch/config.py b/generalresearch/config.py index 76e3995..c6f41e8 100644 --- a/generalresearch/config.py +++ b/generalresearch/config.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from pathlib import Path from pydantic import DirectoryPath, Field, MariaDBDsn, PostgresDsn, RedisDsn diff --git a/generalresearch/grliq/managers/event_plotter.py b/generalresearch/grliq/managers/event_plotter.py index 0bba7c5..94b70ef 100644 --- a/generalresearch/grliq/managers/event_plotter.py +++ b/generalresearch/grliq/managers/event_plotter.py @@ -13,7 +13,7 @@ def make_events_svg( mouse_events: list[MouseEvent], keyboard_events: list[KeyboardEvent] ) -> str: if len(mouse_events) + len(keyboard_events) == 0: - return f'\n' + "\n" + return '\n' + "\n" t = np.array([pm.timeStamp for pm in mouse_events]) t_diff = t.max() - t.min() @@ -88,7 +88,7 @@ def make_events_svg( svg_elements.append(svg_multiline_text(text, cx + 5, cy - 5, font_size)) svg = ( - f'' + '' + "\n".join(svg_elements) + "\n" ) diff --git a/generalresearch/grliq/managers/forensic_data.py b/generalresearch/grliq/managers/forensic_data.py index 7567552..093f7ae 100644 --- a/generalresearch/grliq/managers/forensic_data.py +++ b/generalresearch/grliq/managers/forensic_data.py @@ -1,8 +1,8 @@ from __future__ import annotations +from collections.abc import Collection from datetime import datetime from typing import Any -from collections.abc import Collection from psycopg import sql from pydantic import NonNegativeInt, PositiveInt @@ -611,7 +611,7 @@ class GrlIqDataManager: if res and res["c"] >= 0: return int(res["c"]) - except (Exception,) as e: + except Exception: pass query = f""" diff --git a/generalresearch/grliq/managers/forensic_events.py b/generalresearch/grliq/managers/forensic_events.py index c847d4d..93da481 100644 --- a/generalresearch/grliq/managers/forensic_events.py +++ b/generalresearch/grliq/managers/forensic_events.py @@ -36,36 +36,35 @@ class GrlIqEventManager: "uuid": uuid4().hex, } - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (session_uuid,)) - # Try to update first - update_query = sql.SQL(""" + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (session_uuid,)) + # Try to update first + update_query = sql.SQL(""" UPDATE grliq_forensicevents SET timing_data = %(timing_data)s WHERE session_uuid = %(session_uuid)s AND timing_data IS NULL RETURNING id """) - c.execute(update_query, data) - result = c.fetchone() + c.execute(update_query, data) + result = c.fetchone() - if result: - pk = result["id"] - conn.commit() - return pk + if result: + pk = result["id"] + conn.commit() + return pk - # No matching row to update. Do an insert - insert_query = sql.SQL(""" + # No matching row to update. Do an insert + insert_query = sql.SQL(""" INSERT INTO grliq_forensicevents (uuid, session_uuid, timing_data) VALUES (%(uuid)s, %(session_uuid)s, %(timing_data)s) RETURNING id """) - c.execute(insert_query, data) - pk = c.fetchone()["id"] - conn.commit() + c.execute(insert_query, data) + pk = c.fetchone()["id"] + conn.commit() return int(pk) @@ -88,11 +87,10 @@ class GrlIqEventManager: "event_end": event_end, } - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (session_uuid,)) - # Try to update first - update_query = sql.SQL(""" + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (session_uuid,)) + # Try to update first + update_query = sql.SQL(""" UPDATE grliq_forensicevents SET events = %(events)s, mouse_events = %(mouse_events)s, @@ -102,16 +100,16 @@ class GrlIqEventManager: AND events IS NULL RETURNING id """) - c.execute(update_query, data) - result = c.fetchone() + c.execute(update_query, data) + result = c.fetchone() - if result: - pk = result["id"] - conn.commit() - return pk + if result: + pk = result["id"] + conn.commit() + return pk - # No matching row to update. Do an insert - insert_query = sql.SQL(""" + # No matching row to update. Do an insert + insert_query = sql.SQL(""" INSERT INTO grliq_forensicevents (uuid, session_uuid, events, mouse_events, event_start, event_end) @@ -120,9 +118,9 @@ class GrlIqEventManager: %(event_start)s, %(event_end)s) RETURNING id """) - c.execute(insert_query, data) - pk = c.fetchone()["id"] - conn.commit() + c.execute(insert_query, data) + pk = c.fetchone()["id"] + conn.commit() return int(pk) @@ -167,10 +165,9 @@ class GrlIqEventManager: {filter_str} ORDER BY {order_by} LIMIT {limit} """ - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=query, params=params) - res = c.fetchall() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=query, params=params) + res = c.fetchall() for x in res: if x.get("mouse_events"): @@ -206,10 +203,9 @@ class GrlIqEventManager: AND timing_data IS NOT NULL ORDER BY session_uuid, fe.id DESC; """) - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, params) - res = c.fetchall() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, params) + res = c.fetchall() for x in res: x["timing_data"] = TimingData.model_validate(x["timing_data"]) diff --git a/generalresearch/grliq/managers/forensic_results.py b/generalresearch/grliq/managers/forensic_results.py index 706a7db..158e582 100644 --- a/generalresearch/grliq/managers/forensic_results.py +++ b/generalresearch/grliq/managers/forensic_results.py @@ -90,10 +90,9 @@ class GrlIqCategoryResultsReader: ORDER BY created_at DESC LIMIT {limit} """ - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, params) - res = c.fetchall() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, params) + res = c.fetchall() for x in res: x["client_ip"] = str(x["client_ip"]) diff --git a/generalresearch/grliq/models/custom_types.py b/generalresearch/grliq/models/custom_types.py index 5c7f155..903e230 100644 --- a/generalresearch/grliq/models/custom_types.py +++ b/generalresearch/grliq/models/custom_types.py @@ -1,6 +1,7 @@ -import annotated_types from typing import Annotated +import annotated_types + GrlIqScore = Annotated[int, annotated_types.Ge(0), annotated_types.Le(100)] GrlIqAvgScore = Annotated[float, annotated_types.Ge(0), annotated_types.Le(100)] GrlIqRate = Annotated[float, annotated_types.Ge(0), annotated_types.Le(1)] diff --git a/generalresearch/grliq/models/forensic_data.py b/generalresearch/grliq/models/forensic_data.py index 8d65696..666cb81 100644 --- a/generalresearch/grliq/models/forensic_data.py +++ b/generalresearch/grliq/models/forensic_data.py @@ -779,7 +779,6 @@ class GrlIqData(BaseModel): minutes=90 ), "expired session" - return None def model_dump_sql(self, **kwargs) -> dict[str, Any]: d = dict() diff --git a/generalresearch/grliq/models/forensic_summary.py b/generalresearch/grliq/models/forensic_summary.py index 6f80768..f5b0f25 100644 --- a/generalresearch/grliq/models/forensic_summary.py +++ b/generalresearch/grliq/models/forensic_summary.py @@ -3,7 +3,6 @@ from __future__ import annotations import random from typing import ( Literal, - Optional, Union, get_args, get_origin, @@ -125,7 +124,7 @@ def generate_GrlIqCheckerResultsSummary(): if base_type == GrlIqCheckerResult: if is_opt: fields[f"{field_name}_avg"] = ( - Optional[GrlIqAvgScore], + GrlIqAvgScore | None, Field(default=None, examples=[random.randint(0, 100)]), ) fields[f"{field_name}_pct_none"] = ( diff --git a/generalresearch/grliq/utils.py b/generalresearch/grliq/utils.py index ca8c6a1..bceaa30 100644 --- a/generalresearch/grliq/utils.py +++ b/generalresearch/grliq/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from pathlib import Path from uuid import UUID diff --git a/generalresearch/grpc.py b/generalresearch/grpc.py index 178521e..f1b5611 100644 --- a/generalresearch/grpc.py +++ b/generalresearch/grpc.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from google.protobuf.duration_pb2 import Duration from google.protobuf.timestamp_pb2 import Timestamp diff --git a/generalresearch/incite/base.py b/generalresearch/incite/base.py index 6888ca2..44647fc 100644 --- a/generalresearch/incite/base.py +++ b/generalresearch/incite/base.py @@ -7,8 +7,9 @@ import re import shutil import subprocess import warnings +from collections.abc import Callable, Sequence from concurrent.futures import Future -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from os import R_OK, access, listdir from os.path import isdir from os.path import join as pjoin @@ -17,8 +18,8 @@ from sys import platform from typing import ( TYPE_CHECKING, Any, + Self, ) -from collections.abc import Callable, Sequence from uuid import uuid4 import dask @@ -42,7 +43,6 @@ from pydantic import ( ) from pydantic.json_schema import SkipJsonSchema from sentry_sdk import capture_exception -from typing import Self from generalresearch.config import is_debug from generalresearch.incite.schemas import ( @@ -481,7 +481,7 @@ class CollectionBase(BaseModel): item.cleanup_partials() def clear_tmp_archives(self) -> None: - regex = re.compile(r"\.parquet\.[0-9a-f]{32}", re.I) + regex = re.compile(r"\.parquet\.[0-9a-f]{32}", re.IGNORECASE) for fn in os.listdir(self.archive_path): if regex.search(fn): @@ -554,7 +554,7 @@ class CollectionBase(BaseModel): try: pq.ParquetDataset(highest_version).read().to_pandas() - except (Exception,): + except Exception: # If the most recent version isn't valid, we don't want to # create a symlink to it. # TODO: We could try to be smart and iterate down the most recent diff --git a/generalresearch/incite/collections/__init__.py b/generalresearch/incite/collections/__init__.py index 38749b3..42c3d31 100644 --- a/generalresearch/incite/collections/__init__.py +++ b/generalresearch/incite/collections/__init__.py @@ -214,13 +214,13 @@ class DFCollectionItem(CollectionItemBase): """, params=[start, finish], ) - except (Exception,) as e: + except Exception as e: capture_exception(error=e) LOG.error(f"_from_mysql Exception: {e}") return None if not res: - LOG.warning(f"_from_mysql query returned nothing") + LOG.warning("_from_mysql query returned nothing") # Return an empty df.DataFrame with the correct columns return empty_dataframe_from_schema(coll._schema) @@ -228,7 +228,7 @@ class DFCollectionItem(CollectionItemBase): df = self.validate_df(df=df) if df is None: - LOG.warning(f"_from_mysql query results failed validation") + LOG.warning("_from_mysql query results failed validation") # Schema validation can fail... return None @@ -265,13 +265,13 @@ class DFCollectionItem(CollectionItemBase): """, params=[start, finish], ) - except (Exception,) as e: + except Exception as e: capture_exception(error=e) LOG.error(f"_from_postgres Exception: {e}") return None if not res: - LOG.warning(f"_from_postgres query returned nothing") + LOG.warning("_from_postgres query returned nothing") # Return an empty df.DataFrame with the correct columns return empty_dataframe_from_schema(coll._schema) @@ -279,7 +279,7 @@ class DFCollectionItem(CollectionItemBase): df = self.validate_df(df=df) if df is None: - LOG.warning(f"_from_postgres query results failed validation") + LOG.warning("_from_postgres query results failed validation") # Schema validation can fail... return None @@ -351,7 +351,7 @@ class DFCollectionItem(CollectionItemBase): c: Cursor = conn.cursor() for chunk in chunked(tx_ids, n=5_000): c.execute( - query=f""" + query=""" SELECT ltm.transaction_id AS tx_id, ltm.id AS tx_metadata_id, ltm.key, ltm.value @@ -466,7 +466,7 @@ class DFCollectionItem(CollectionItemBase): compression="brotli", ) - except (Exception,) as e: + except Exception as e: LOG.exception(e) self.delete_archive(tmp_path) return False @@ -553,7 +553,7 @@ class DFCollectionItem(CollectionItemBase): write_metadata_file=True, compression="brotli", ) - except (Exception,) as e: + except Exception as e: LOG.exception(e) self.delete_archive(next_numbered_path) return False diff --git a/generalresearch/incite/defaults.py b/generalresearch/incite/defaults.py index 5a95607..d4025fc 100644 --- a/generalresearch/incite/defaults.py +++ b/generalresearch/incite/defaults.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from generalresearch.incite.base import GRLDatasets from generalresearch.incite.collections import DFCollectionType diff --git a/generalresearch/incite/mergers/foundations/__init__.py b/generalresearch/incite/mergers/foundations/__init__.py index f7a45a8..d3100fa 100644 --- a/generalresearch/incite/mergers/foundations/__init__.py +++ b/generalresearch/incite/mergers/foundations/__init__.py @@ -117,7 +117,7 @@ def annotate_product_and_team_id( try: with conn.cursor() as c: c.execute( - query=f""" + query=""" SELECT u.id AS user_id, u.product_id, bp.team_id FROM thl_user u diff --git a/generalresearch/incite/mergers/foundations/enriched_session.py b/generalresearch/incite/mergers/foundations/enriched_session.py index a368e6c..8a707ca 100644 --- a/generalresearch/incite/mergers/foundations/enriched_session.py +++ b/generalresearch/incite/mergers/foundations/enriched_session.py @@ -60,17 +60,17 @@ class EnrichedSessionMergeItem(MergeCollectionItem): return # --- Session --- - LOG.warning(f"EnrichedSessionMergeItem: get session_collection") + LOG.warning("EnrichedSessionMergeItem: get session_collection") session_items = [w for w in session_coll.items if w.interval.overlaps(ir)] if len(session_items) == 0: - LOG.warning(f"EnrichedSessionMergeItem: no session items. set_empty.") + LOG.warning("EnrichedSessionMergeItem: no session items. set_empty.") if self.should_archive(): self.set_empty() return if not ( session_items[-1].has_partial_archive() or session_items[-1].has_archive() ): - LOG.warning(f"EnrichedSessionMergeItem: session isn't updated!") + LOG.warning("EnrichedSessionMergeItem: session isn't updated!") return sddf = session_coll.ddf( @@ -81,7 +81,7 @@ class EnrichedSessionMergeItem(MergeCollectionItem): ) # --- Walls --- - LOG.warning(f"EnrichedSessionMergeItem: merge wall_collection") + LOG.warning("EnrichedSessionMergeItem: merge wall_collection") wall_items = [ w for w in wall_coll.items @@ -95,7 +95,7 @@ class EnrichedSessionMergeItem(MergeCollectionItem): ] if len(wall_items) == 0: - LOG.error(f"EnrichedSessionMergeItem: no wall items") + LOG.error("EnrichedSessionMergeItem: no wall items") return wddf = wall_coll.ddf( diff --git a/generalresearch/incite/mergers/foundations/enriched_task_adjust.py b/generalresearch/incite/mergers/foundations/enriched_task_adjust.py index f3ab8d8..234cd5b 100644 --- a/generalresearch/incite/mergers/foundations/enriched_task_adjust.py +++ b/generalresearch/incite/mergers/foundations/enriched_task_adjust.py @@ -55,7 +55,7 @@ class EnrichedTaskAdjustMergeItem(MergeCollectionItem): LOG.warning(f"EnrichedReconMergeItem.build({ir})") # --- Task Adjustments --- - LOG.warning(f"EnrichedReconMergeItem: get session_collection") + LOG.warning("EnrichedReconMergeItem: get session_collection") task_adj_coll_items = [ w for w in task_adj_coll.items if w.interval.overlaps(ir) ] @@ -209,6 +209,5 @@ class EnrichedTaskAdjustMerge(MergeCollection): enriched_wall=enriched_wall, pg_config=pg_config, ) - except (Exception,) as e: + except Exception as e: capture_exception(error=e) - pass diff --git a/generalresearch/incite/mergers/foundations/enriched_wall.py b/generalresearch/incite/mergers/foundations/enriched_wall.py index 5a7dd2b..b2ac7bb 100644 --- a/generalresearch/incite/mergers/foundations/enriched_wall.py +++ b/generalresearch/incite/mergers/foundations/enriched_wall.py @@ -55,10 +55,10 @@ class EnrichedWallMergeItem(MergeCollectionItem): return # --- Wall --- - LOG.warning(f"EnrichedWallMergeItem: get wall_collection") + LOG.warning("EnrichedWallMergeItem: get wall_collection") wall_items = [w for w in wall_coll.items if w.interval.overlaps(ir)] if len(wall_items) == 0: - LOG.warning(f"EnrichedWallMergeItem: no wall items. set_empty.") + LOG.warning("EnrichedWallMergeItem: no wall items. set_empty.") if self.should_archive(): self.set_empty() return @@ -93,7 +93,7 @@ class EnrichedWallMergeItem(MergeCollectionItem): wdf = wdf.reset_index(drop=False) # --- Sessions --- - LOG.warning(f"EnrichedWallMergeItem: merge session_collection") + LOG.warning("EnrichedWallMergeItem: merge session_collection") session_items = [ s for s in session_coll.items @@ -107,7 +107,7 @@ class EnrichedWallMergeItem(MergeCollectionItem): ] if len(session_items) == 0: - LOG.error(f"EnrichedWallMergeItem: no session items. breaking early.") + LOG.error("EnrichedWallMergeItem: no session items. breaking early.") return sdf = session_coll.ddf( diff --git a/generalresearch/incite/mergers/ym_survey_wall.py b/generalresearch/incite/mergers/ym_survey_wall.py index c060aae..9750b57 100644 --- a/generalresearch/incite/mergers/ym_survey_wall.py +++ b/generalresearch/incite/mergers/ym_survey_wall.py @@ -62,7 +62,7 @@ class YMSurveyWallMergeCollectionItem(MergeCollectionItem): ) ddf = ddf[ddf["started"] > start] - LOG.warning(f"YMSurveyWallMerge: merge session_collection") + LOG.warning("YMSurveyWallMerge: merge session_collection") session_items = [ s for s in enriched_session.items @@ -98,17 +98,16 @@ class YMSurveyWallMergeCollectionItem(MergeCollectionItem): df.dropna(subset="product_id", how="any", inplace=True) df.sort_values(by="started", inplace=True) - LOG.debug(f"YMSurveyWallMerge.build() validation") + LOG.debug("YMSurveyWallMerge.build() validation") df = self.validate_df(df=df) if df is not None: ddf = dd.from_pandas(df, npartitions=4) - LOG.info(f"YMSurveyWallMerge.build() saving") + LOG.info("YMSurveyWallMerge.build() saving") self.to_archive_symlink(client=client, ddf=ddf) else: LOG.warning("YMSurveyWallMerge failed validation") - return None class YMSurveyWallMerge(MergeCollection): @@ -144,8 +143,7 @@ class YMSurveyWallMerge(MergeCollection): wall_coll=wall_coll, enriched_session=enriched_session, ) - except (Exception,) as e: + except Exception as e: capture_exception(error=e) - pass item.delete_dangling_partials(keep_latest=2, target_path=item.path) diff --git a/generalresearch/incite/mergers/ym_wall_summary.py b/generalresearch/incite/mergers/ym_wall_summary.py index 618b810..37fc3b9 100644 --- a/generalresearch/incite/mergers/ym_wall_summary.py +++ b/generalresearch/incite/mergers/ym_wall_summary.py @@ -117,9 +117,8 @@ class YMWallSummaryMerge(MergeCollection): # item every time build is run even if it isn't closed # if item.should_archive(): item.fetch(wall_collection, session_collection, user_id_product) - except (Exception,) as e: + except Exception as e: capture_exception(e) - pass @staticmethod def build_groupbys(df: pd.DataFrame) -> pd.DataFrame: diff --git a/generalresearch/incite/schemas/thl_web.py b/generalresearch/incite/schemas/thl_web.py index b831b9a..c1be202 100644 --- a/generalresearch/incite/schemas/thl_web.py +++ b/generalresearch/incite/schemas/thl_web.py @@ -1,4 +1,4 @@ -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index, MultiIndex diff --git a/generalresearch/locales/setup_json.py b/generalresearch/locales/setup_json.py index 356084a..57caa00 100644 --- a/generalresearch/locales/setup_json.py +++ b/generalresearch/locales/setup_json.py @@ -11,6 +11,7 @@ def country_default_lang(): """ raise ValueError("no need to run this, I already ran it.") import pandas as pd + from generalresearch.locales import Localelator l = Localelator() diff --git a/generalresearch/locales/timezone.py b/generalresearch/locales/timezone.py index fce6e0e..810dba3 100644 --- a/generalresearch/locales/timezone.py +++ b/generalresearch/locales/timezone.py @@ -1,4 +1,3 @@ -from typing import Optional from pytz import country_timezones diff --git a/generalresearch/managers/cint/survey.py b/generalresearch/managers/cint/survey.py index da1ecd9..819ae3d 100644 --- a/generalresearch/managers/cint/survey.py +++ b/generalresearch/managers/cint/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/criteria.py b/generalresearch/managers/criteria.py index b5d9830..c13b8ac 100644 --- a/generalresearch/managers/criteria.py +++ b/generalresearch/managers/criteria.py @@ -2,7 +2,7 @@ from __future__ import annotations from abc import ABC from collections.abc import Collection -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from more_itertools import chunked @@ -30,7 +30,6 @@ class CriteriaManager(SqlManager, ABC): """ Create a single criterion """ - ... def filter(self, hashes: Collection[str]) -> dict[str, MarketplaceCondition]: """ @@ -96,7 +95,6 @@ class CriteriaManager(SqlManager, ABC): ) conn.commit() - return None @property def mysql_fields(self) -> str: diff --git a/generalresearch/managers/dynata/survey.py b/generalresearch/managers/dynata/survey.py index 3a15c4d..7643dc4 100644 --- a/generalresearch/managers/dynata/survey.py +++ b/generalresearch/managers/dynata/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/events.py b/generalresearch/managers/events.py index f3c6a04..f6c429e 100644 --- a/generalresearch/managers/events.py +++ b/generalresearch/managers/events.py @@ -5,7 +5,7 @@ import math import socket import threading import time -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import TYPE_CHECKING @@ -724,7 +724,6 @@ class EventManager(StatsManager): ) ) self.publish_event(msg, product_id=user.product_id) - return def handle_task_finish(self, wall: Wall, session: Session, user: User): self.mark_user_active(user=user) @@ -818,7 +817,6 @@ class EventSubscriber(RedisManager): p.subscribe(self.get_channel_name()) self.pubsub_client = r self.pubsub = p - return def get_channel_name(self): return f"{self.cache_prefix}:event-channel:{self.product_id}" diff --git a/generalresearch/managers/gr/authentication.py b/generalresearch/managers/gr/authentication.py index 409cb10..80bee4b 100644 --- a/generalresearch/managers/gr/authentication.py +++ b/generalresearch/managers/gr/authentication.py @@ -3,7 +3,7 @@ from __future__ import annotations import binascii import logging import os -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any from psycopg import sql @@ -146,7 +146,7 @@ class GRUserManager(PostgresManagerWithRedis): for item in res: for k, v in item.items(): - if isinstance(item[k], datetime): + if isinstance(v, datetime): item[k] = item[k].replace(tzinfo=UTC) return [GRUser.model_validate(item) for item in res] @@ -270,7 +270,6 @@ class GRTokenManager(PostgresManager): ) conn.commit() - return def get_by_user_id(self, user_id: PositiveInt) -> GRToken | None: # django authtoken_token table has (user_id) UNIQUE constraint diff --git a/generalresearch/managers/gr/team.py b/generalresearch/managers/gr/team.py index 393f446..d04370b 100644 --- a/generalresearch/managers/gr/team.py +++ b/generalresearch/managers/gr/team.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from typing import TYPE_CHECKING from uuid import uuid4 diff --git a/generalresearch/managers/innovate/survey.py b/generalresearch/managers/innovate/survey.py index c65b100..cddfba2 100644 --- a/generalresearch/managers/innovate/survey.py +++ b/generalresearch/managers/innovate/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/leaderboard/__init__.py b/generalresearch/managers/leaderboard/__init__.py index d5138cd..048d2cb 100644 --- a/generalresearch/managers/leaderboard/__init__.py +++ b/generalresearch/managers/leaderboard/__init__.py @@ -1,8 +1,9 @@ from __future__ import annotations +from zoneinfo import ZoneInfo + import pytz from cachetools import LRUCache, cached -from zoneinfo import ZoneInfo @cached(cache=LRUCache(maxsize=1)) diff --git a/generalresearch/managers/leaderboard/manager.py b/generalresearch/managers/leaderboard/manager.py index 71c3a73..07e3e2c 100644 --- a/generalresearch/managers/leaderboard/manager.py +++ b/generalresearch/managers/leaderboard/manager.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from decimal import Decimal from functools import cached_property from typing import TYPE_CHECKING diff --git a/generalresearch/managers/morning/survey.py b/generalresearch/managers/morning/survey.py index 5fba70d..0e29010 100644 --- a/generalresearch/managers/morning/survey.py +++ b/generalresearch/managers/morning/survey.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/network/label.py b/generalresearch/managers/network/label.py index c5306a6..1f44862 100644 --- a/generalresearch/managers/network/label.py +++ b/generalresearch/managers/network/label.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from psycopg import sql from pydantic import IPvAnyNetwork, TypeAdapter @@ -28,10 +28,9 @@ class IPLabelManager(PostgresManager): %(provider)s, %(metadata)s ) RETURNING id;""") params = ip_label.model_dump_postgres() - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, params) - pk = c.fetchone()["id"] + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, params) + pk = c.fetchone()["id"] return ip_label def make_filter_str( diff --git a/generalresearch/managers/network/tool_run.py b/generalresearch/managers/network/tool_run.py index 026b3d3..73afc13 100644 --- a/generalresearch/managers/network/tool_run.py +++ b/generalresearch/managers/network/tool_run.py @@ -49,7 +49,6 @@ class ToolRunManager(PostgresManager): c.execute(query, params) run_id = c.fetchone()["id"] run.id = run_id - return None def create_tool_run(self, run: NmapRun | RDNSRun | MTRRun): if type(run) is NmapRun: @@ -77,10 +76,9 @@ class ToolRunManager(PostgresManager): """ Insert a PortScan + PortScanPorts from a Pydantic NmapResult. """ - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - self._create_tool_run(run, c) - self.nmap_manager._create(run, c=c) + with self.pg_config.make_connection() as conn, conn.cursor() as c: + self._create_tool_run(run, c) + self.nmap_manager._create(run, c=c) return run def get_nmap_run(self, id: int) -> NmapRun: @@ -98,10 +96,9 @@ class ToolRunManager(PostgresManager): """ Insert a RDnsRun + RDNSResult """ - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - self._create_tool_run(run, c) - self.rdns_manager._create(run, c=c) + with self.pg_config.make_connection() as conn, conn.cursor() as c: + self._create_tool_run(run, c) + self.rdns_manager._create(run, c=c) return run def get_rdns_run(self, id: int) -> RDNSRun: @@ -120,10 +117,9 @@ class ToolRunManager(PostgresManager): return RDNSRun.model_validate(res) def create_mtr_run(self, run: MTRRun) -> MTRRun: - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - self._create_tool_run(run, c) - self.mtr_manager._create(run, c=c) + with self.pg_config.make_connection() as conn, conn.cursor() as c: + self._create_tool_run(run, c) + self.mtr_manager._create(run, c=c) return run def get_mtr_run(self, id: int) -> MTRRun: diff --git a/generalresearch/managers/precision/survey.py b/generalresearch/managers/precision/survey.py index 833cb28..c13dca8 100644 --- a/generalresearch/managers/precision/survey.py +++ b/generalresearch/managers/precision/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError @@ -125,7 +125,7 @@ class PrecisionSurveyManager(SurveyManager): country_data = [(survey.survey_id, c) for c in survey.country_isos] c.executemany( - f""" + """ INSERT INTO `thl-precision`.`precision_survey_country` (survey_id, country_iso, is_active) VALUES (%s, %s, TRUE) @@ -134,7 +134,7 @@ class PrecisionSurveyManager(SurveyManager): ) lang_data = [(survey.survey_id, c) for c in survey.language_isos] c.executemany( - f""" + """ INSERT INTO `thl-precision`.`precision_survey_language` (survey_id, language_iso, is_active) VALUES (%s, %s, TRUE) @@ -188,7 +188,7 @@ class PrecisionSurveyManager(SurveyManager): country_data = [(survey.survey_id, c) for c in survey.country_isos] # Turn ON countries in this survey's list of countries, insert row, if already exists, set active. c.executemany( - query=f""" + query=""" INSERT INTO `thl-precision`.`precision_survey_country` (survey_id, country_iso, is_active) VALUES (%s, %s, TRUE) ON DUPLICATE KEY UPDATE is_active = TRUE; @@ -207,7 +207,7 @@ class PrecisionSurveyManager(SurveyManager): ) language_data = [(survey.survey_id, c) for c in survey.language_isos] c.executemany( - query=f""" + query=""" INSERT INTO `thl-precision`.`precision_survey_language` (survey_id, language_iso, is_active) VALUES (%s, %s, TRUE) ON DUPLICATE KEY UPDATE is_active = TRUE; diff --git a/generalresearch/managers/prodege/survey.py b/generalresearch/managers/prodege/survey.py index 750383f..983572c 100644 --- a/generalresearch/managers/prodege/survey.py +++ b/generalresearch/managers/prodege/survey.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql diff --git a/generalresearch/managers/repdata/survey.py b/generalresearch/managers/repdata/survey.py index 1e1c3c6..fe6f621 100644 --- a/generalresearch/managers/repdata/survey.py +++ b/generalresearch/managers/repdata/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql @@ -105,7 +105,7 @@ class RepDataSurveyManager(SurveyManager): surveys = {s.survey_id: s for s in surveys} if surveys: res = self.sql_helper.execute_sql_query( - query=f""" + query=""" SELECT * FROM `thl-repdata`.`repdata_surveystream` WHERE survey_id IN %s diff --git a/generalresearch/managers/sago/survey.py b/generalresearch/managers/sago/survey.py index 2582902..462d2ef 100644 --- a/generalresearch/managers/sago/survey.py +++ b/generalresearch/managers/sago/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/spectrum/survey.py b/generalresearch/managers/spectrum/survey.py index 58f8a1a..9b58d43 100644 --- a/generalresearch/managers/spectrum/survey.py +++ b/generalresearch/managers/spectrum/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime import pymysql from pymysql import IntegrityError diff --git a/generalresearch/managers/survey.py b/generalresearch/managers/survey.py index 964344f..1f057fb 100644 --- a/generalresearch/managers/survey.py +++ b/generalresearch/managers/survey.py @@ -12,14 +12,12 @@ class SurveyManager(SqlManager, ABC): """ Create a single survey """ - ... def update(self, surveys: list[MarketplaceTask]) -> bool: """ Update a list of surveys. Depending on the implementation, this may operate one by one or as a bulk update. """ - ... def update_field(self, survey: MarketplaceTask, field: str) -> bool: """ diff --git a/generalresearch/managers/thl/buyer.py b/generalresearch/managers/thl/buyer.py index ae40bc8..2cb582f 100644 --- a/generalresearch/managers/thl/buyer.py +++ b/generalresearch/managers/thl/buyer.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from generalresearch.managers.base import Permission, PostgresManager from generalresearch.models import Source diff --git a/generalresearch/managers/thl/cashout_method.py b/generalresearch/managers/thl/cashout_method.py index 94617f7..10282d4 100644 --- a/generalresearch/managers/thl/cashout_method.py +++ b/generalresearch/managers/thl/cashout_method.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Collection from copy import copy -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from typing import Any from uuid import UUID, uuid4 diff --git a/generalresearch/managers/thl/contest_manager.py b/generalresearch/managers/thl/contest_manager.py index 286de3d..62146d7 100644 --- a/generalresearch/managers/thl/contest_manager.py +++ b/generalresearch/managers/thl/contest_manager.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from typing import Any, Literal, cast from uuid import UUID @@ -827,7 +827,6 @@ class MilestoneContestManager(ContestBaseManager): ) self.end_milestone_contest(contest) - return None def enter_contest_db_work_milestone( self, contest: MilestoneUserView, user: User, incr: PositiveInt diff --git a/generalresearch/managers/thl/ledger_manager/conditions.py b/generalresearch/managers/thl/ledger_manager/conditions.py index f79a0f0..b2fd465 100644 --- a/generalresearch/managers/thl/ledger_manager/conditions.py +++ b/generalresearch/managers/thl/ledger_manager/conditions.py @@ -1,9 +1,9 @@ from __future__ import annotations import logging -from datetime import datetime, timedelta, timezone, UTC -from typing import TYPE_CHECKING from collections.abc import Callable +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING from generalresearch.config import JAMES_BILLINGS_BPID, JAMES_BILLINGS_TX_CUTOFF from generalresearch.currency import USDCent diff --git a/generalresearch/managers/thl/ledger_manager/exceptions.py b/generalresearch/managers/thl/ledger_manager/exceptions.py index b79c153..48c102b 100644 --- a/generalresearch/managers/thl/ledger_manager/exceptions.py +++ b/generalresearch/managers/thl/ledger_manager/exceptions.py @@ -11,7 +11,6 @@ class LedgerTransactionCreateError(Exception): Ledger transaction creation failed """ - pass class LedgerTransactionCreateLockError(LedgerTransactionCreateError): @@ -19,7 +18,6 @@ class LedgerTransactionCreateLockError(LedgerTransactionCreateError): Ledger transaction creation failed because we could not acquire a lock """ - pass class LedgerTransactionReleaseLockError(LedgerTransactionCreateError): @@ -29,7 +27,6 @@ class LedgerTransactionReleaseLockError(LedgerTransactionCreateError): back-populate as in sentry I see this very rarely. """ - pass class LedgerTransactionFlagAlreadyExistsError(LedgerTransactionCreateError): @@ -38,7 +35,6 @@ class LedgerTransactionFlagAlreadyExistsError(LedgerTransactionCreateError): tx was already set """ - pass class LedgerTransactionConditionFailedError(LedgerTransactionCreateError): @@ -46,4 +42,3 @@ class LedgerTransactionConditionFailedError(LedgerTransactionCreateError): We tried to create a transaction but the condition check failed. """ - pass diff --git a/generalresearch/managers/thl/ledger_manager/ledger.py b/generalresearch/managers/thl/ledger_manager/ledger.py index 00fac27..410f6ca 100644 --- a/generalresearch/managers/thl/ledger_manager/ledger.py +++ b/generalresearch/managers/thl/ledger_manager/ledger.py @@ -2,10 +2,9 @@ from __future__ import annotations import logging from collections import defaultdict -from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from collections.abc import Callable, Collection +from datetime import UTC, datetime, timedelta from typing import Any -from collections.abc import Callable from uuid import UUID import redis @@ -343,7 +342,7 @@ class LedgerTransactionManager(LedgerManagerBasePostgres): assert len(tag) > 6, "Please confirm the tag is valid" res = self.pg_config.execute_sql_query( - query=f""" + query=""" SELECT lt.id FROM ledger_transaction AS lt WHERE tag = %s @@ -361,7 +360,7 @@ class LedgerTransactionManager(LedgerManagerBasePostgres): def get_tx_ids_by_tags(self, tags: list[str]) -> set[PositiveInt]: res = self.pg_config.execute_sql_query( - query=f""" + query=""" SELECT lt.id, lt.tag, lt.created, lt.ext_description FROM ledger_transaction AS lt WHERE tag = ANY(%s) @@ -869,7 +868,7 @@ class LedgerAccountManager(LedgerManagerBasePostgres): # qualified_name has a unique index so there can only be 0 or 1 match. res = self.pg_config.execute_sql_query( - query=f""" + query=""" SELECT uuid, display_name, qualified_name, account_type, normal_balance, reference_type, @@ -928,7 +927,7 @@ class LedgerAccountManager(LedgerManagerBasePostgres): # TODO: Move to RR with long timeout (2min+), it causes problems res = self.pg_config.execute_sql_query( - query=f""" + query=""" SELECT SUM(amount * direction) AS total FROM ledger_entry WHERE account_id = %s @@ -1046,7 +1045,7 @@ class LedgerManager( """This is for testing only, as it'll take forever to run this if the ledger_manager is huge """ - res = self.pg_config.execute_sql_query(f""" + res = self.pg_config.execute_sql_query(""" SELECT SUM(CASE WHEN normal_balance = -1 THEN total ELSE 0 END) AS credit_total, SUM(CASE WHEN normal_balance = 1 THEN total ELSE 0 END) AS debit_total diff --git a/generalresearch/managers/thl/ledger_manager/thl_ledger.py b/generalresearch/managers/thl/ledger_manager/thl_ledger.py index c9203f6..ded6518 100644 --- a/generalresearch/managers/thl/ledger_manager/thl_ledger.py +++ b/generalresearch/managers/thl/ledger_manager/thl_ledger.py @@ -1,11 +1,10 @@ from __future__ import annotations import logging -from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from collections.abc import Callable, Collection +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import TYPE_CHECKING -from collections.abc import Callable from uuid import UUID import numpy as np @@ -407,8 +406,7 @@ class ThlLedgerManager(LedgerManager): f"bp_pay {bp_pay} > thl_net {thl_net}. Capping bp_pay to thl_net." ) bp_pay = thl_net - if user_pay > bp_pay: - user_pay = bp_pay + user_pay = min(user_pay, bp_pay) commission_amount = round(thl_net - bp_pay) @@ -538,7 +536,7 @@ class ThlLedgerManager(LedgerManager): ] else: - logger.info(f"create_transaction_task_adjustment. No transactions needed.") + logger.info("create_transaction_task_adjustment. No transactions needed.") return None amt_str = f"${abs(change_amount) / 100:,.2f}" @@ -666,7 +664,7 @@ class ThlLedgerManager(LedgerManager): else: logger.info( - f"create_transaction_bp_adjustment. No transactions needed." + "create_transaction_bp_adjustment. No transactions needed." ) return None else: @@ -738,7 +736,7 @@ class ThlLedgerManager(LedgerManager): else: logger.info( - f"create_transaction_bp_adjustment. No transactions needed." + "create_transaction_bp_adjustment. No transactions needed." ) return None @@ -856,7 +854,7 @@ class ThlLedgerManager(LedgerManager): ), ] - ext_description = f"BP Payout" + ext_description = "BP Payout" t = self.create_tx( entries=entries, metadata=metadata, @@ -983,7 +981,7 @@ class ThlLedgerManager(LedgerManager): raise ValueError("Invalid Direction") if description is None: - description = f"BP Plug" + description = "BP Plug" t = self.create_tx( entries=entries, @@ -1191,7 +1189,7 @@ class ThlLedgerManager(LedgerManager): f"Trying to cancel user payout {payout_event.uuid} with no request tx found." ) - description = f"User Payout Cancelled" + description = "User Payout Cancelled" f = lambda: self.create_tx_user_payout_cancelled_( user=user, payout_event=payout_event, @@ -1900,7 +1898,7 @@ class ThlLedgerManager(LedgerManager): reserve = round(wall["user_payout_int"].sum() - wall["redeemable"].sum()) redeemable_balance = user_wallet_balance - reserve - redeemable_balance = 0 if redeemable_balance < 0 else redeemable_balance + redeemable_balance = max(redeemable_balance, 0) if redeemable_balance > 0: # it is possible the user_wallet_balance is negative, in which case diff --git a/generalresearch/managers/thl/product.py b/generalresearch/managers/thl/product.py index 3b92361..d924e17 100644 --- a/generalresearch/managers/thl/product.py +++ b/generalresearch/managers/thl/product.py @@ -4,7 +4,7 @@ import json import logging import operator from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from decimal import Decimal from threading import Lock from typing import TYPE_CHECKING @@ -94,9 +94,9 @@ class ProductManager(PostgresManager): return self.fetch_uuids( product_uuids=[product_uuid], )[0] - except (AssertionError,): + except AssertionError: return None - except (IndexError,): + except IndexError: return None def get_by_uuids_if_exists( @@ -397,11 +397,11 @@ class ProductManager(PostgresManager): # # from pymysql import IntegrityError # except IntegrityError as e: - except (Exception,) as e: + except Exception as e: try: return self.get_by_uuid(product_uuid=instance.id) - except (Exception,) as e2: + except Exception: pass finally: self.cache_clear(instance.id) diff --git a/generalresearch/managers/thl/profiling/question.py b/generalresearch/managers/thl/profiling/question.py index 10a9e32..078894a 100644 --- a/generalresearch/managers/thl/profiling/question.py +++ b/generalresearch/managers/thl/profiling/question.py @@ -85,7 +85,7 @@ class QuestionManager(PostgresManager): def lookup_by_property( self, property_code: str, country_iso: str, language_iso: str ) -> UpkQuestion: - query = f""" + query = """ SELECT data, property_code, explanation_template, explanation_fragment_template FROM marketplace_question WHERE property_code = %(property_code)s diff --git a/generalresearch/managers/thl/profiling/uqa.py b/generalresearch/managers/thl/profiling/uqa.py index 3854333..fa1747b 100644 --- a/generalresearch/managers/thl/profiling/uqa.py +++ b/generalresearch/managers/thl/profiling/uqa.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from generalresearch.managers.base import PostgresManagerWithRedis from generalresearch.models.thl.profiling.user_question_answer import ( diff --git a/generalresearch/managers/thl/profiling/user_upk.py b/generalresearch/managers/thl/profiling/user_upk.py index ee36124..a2cddb3 100644 --- a/generalresearch/managers/thl/profiling/user_upk.py +++ b/generalresearch/managers/thl/profiling/user_upk.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from collections import defaultdict from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from typing import Any from uuid import UUID diff --git a/generalresearch/managers/thl/session.py b/generalresearch/managers/thl/session.py index 771f882..959003c 100644 --- a/generalresearch/managers/thl/session.py +++ b/generalresearch/managers/thl/session.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Collection -from datetime import datetime, timedelta, timezone, UTC +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import Any from uuid import UUID, uuid4 @@ -190,14 +190,7 @@ class SessionManager(PostgresManager): # re-run model_validate after finished = finished if finished else datetime.now(tz=UTC) session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "status_code_2": status_code_2, - "finished": finished, - "payout": payout, - "user_payout": user_payout, - } + status=status, status_code_1=status_code_1, status_code_2=status_code_2, finished=finished, payout=payout, user_payout=user_payout ) d = session.model_dump_mysql() self.pg_config.execute_write( diff --git a/generalresearch/managers/thl/survey.py b/generalresearch/managers/thl/survey.py index 966e96d..a9ec841 100644 --- a/generalresearch/managers/thl/survey.py +++ b/generalresearch/managers/thl/survey.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections import defaultdict from collections.abc import Collection -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any import pandas as pd @@ -280,7 +280,6 @@ class SurveyManager(PostgresManager): query=query, params={"survey_pks": survey_pks}, ) - return None def update_surveys_categories(self, surveys: list[Survey] | None = None) -> None: for chunk in chunked(surveys, 500): @@ -329,12 +328,11 @@ class SurveyManager(PostgresManager): ] with self.pg_config.make_connection() as conn: # noinspection PyArgumentList - with conn.transaction(): - with conn.cursor() as c: - c.execute(temp_table_sql) - c.executemany(insert_values_sql, rows) - c.execute(delete_sql) - c.execute(upsert_sql) + with conn.transaction(), conn.cursor() as c: + c.execute(temp_table_sql) + c.executemany(insert_values_sql, rows) + c.execute(delete_sql) + c.execute(upsert_sql) conn.commit() def get_survey_categories(self): @@ -760,12 +758,11 @@ class SurveyStatManager(PostgresManager): print(query) print(params) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute("SET work_mem = '256MB';") - c.execute("SET statement_timeout = '10s';") - c.execute(query, params=params) - res = c.fetchall() + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute("SET work_mem = '256MB';") + c.execute("SET statement_timeout = '10s';") + c.execute(query, params=params) + res = c.fetchall() return [SurveyStat.model_validate(x) for x in res] diff --git a/generalresearch/managers/thl/task_adjustment.py b/generalresearch/managers/thl/task_adjustment.py index 3ec3d41..60bade7 100644 --- a/generalresearch/managers/thl/task_adjustment.py +++ b/generalresearch/managers/thl/task_adjustment.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from functools import cached_property @@ -148,10 +148,7 @@ class TaskAdjustmentManager(PostgresManager): if ( wall.status == Status.COMPLETE and adjusted_status == WallAdjustedStatus.ADJUSTED_TO_COMPLETE - ): - new_adjusted_status = None - new_adjusted_cpi = None - elif ( + ) or ( wall.status != Status.COMPLETE and adjusted_status == WallAdjustedStatus.ADJUSTED_TO_FAIL ): diff --git a/generalresearch/managers/thl/user_compensate.py b/generalresearch/managers/thl/user_compensate.py index c6c0747..8338424 100644 --- a/generalresearch/managers/thl/user_compensate.py +++ b/generalresearch/managers/thl/user_compensate.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from uuid import uuid4 diff --git a/generalresearch/managers/thl/user_manager/mysql_user_manager.py b/generalresearch/managers/thl/user_manager/mysql_user_manager.py index 7931ba4..d2d0ffc 100644 --- a/generalresearch/managers/thl/user_manager/mysql_user_manager.py +++ b/generalresearch/managers/thl/user_manager/mysql_user_manager.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from collections.abc import Collection -from datetime import datetime, timezone, UTC +from datetime import UTC, datetime from functools import lru_cache from uuid import uuid4 diff --git a/generalresearch/managers/thl/user_manager/user_manager.py b/generalresearch/managers/thl/user_manager/user_manager.py index 3794020..7486869 100644 --- a/generalresearch/managers/thl/user_manager/user_manager.py +++ b/generalresearch/managers/thl/user_manager/user_manager.py @@ -103,7 +103,6 @@ class UserManager: event_value=event_value, ) - return None def cache_clear(self): # Generally this is used in testing. This clears the .get_user's lru_cache. diff --git a/generalresearch/managers/thl/userhealth.py b/generalresearch/managers/thl/userhealth.py index 0bc60ec..babed04 100644 --- a/generalresearch/managers/thl/userhealth.py +++ b/generalresearch/managers/thl/userhealth.py @@ -2,7 +2,7 @@ from __future__ import annotations import ipaddress from collections.abc import Collection -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import zip_longest from typing import Any diff --git a/generalresearch/managers/thl/wall.py b/generalresearch/managers/thl/wall.py index 7e413d7..03ca1c6 100644 --- a/generalresearch/managers/thl/wall.py +++ b/generalresearch/managers/thl/wall.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging from collections import defaultdict from collections.abc import Collection -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from functools import cached_property from uuid import uuid4 diff --git a/generalresearch/managers/thl/wallet/tango.py b/generalresearch/managers/thl/wallet/tango.py index a4ebf22..4abfc70 100644 --- a/generalresearch/managers/thl/wallet/tango.py +++ b/generalresearch/managers/thl/wallet/tango.py @@ -44,7 +44,7 @@ def complete_tango_order( tango_client=tango_client, ) - except Exception as e: + except Exception: # todo: its possible the order went through, but something else was wrong # we should try to retrieve the order by its ref_id and confirm it really # failed... diff --git a/generalresearch/models/admin/request.py b/generalresearch/models/admin/request.py index 2d68de1..5fdc784 100644 --- a/generalresearch/models/admin/request.py +++ b/generalresearch/models/admin/request.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from enum import Enum from typing import Literal diff --git a/generalresearch/models/custom_types.py b/generalresearch/models/custom_types.py index 9346064..c200b34 100644 --- a/generalresearch/models/custom_types.py +++ b/generalresearch/models/custom_types.py @@ -2,8 +2,7 @@ from __future__ import annotations import json import re -import sys as _sys -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from typing import Annotated, Any, Literal from uuid import UUID diff --git a/generalresearch/models/dynata/survey.py b/generalresearch/models/dynata/survey.py index 9adac42..5a9f763 100644 --- a/generalresearch/models/dynata/survey.py +++ b/generalresearch/models/dynata/survey.py @@ -132,9 +132,7 @@ class DynataCondition(MarketplaceCondition): if cell["kind"] == "RANGE": d["values"] = [ - "{0}-{1}".format( - cell["range"]["from"] or "inf", cell["range"]["to"] or "inf" - ) + f"{cell["range"]["from"] or "inf"}-{cell["range"]["to"] or "inf"}" ] d["value_type"] = ConditionValueType.RANGE return cls.model_validate(d) diff --git a/generalresearch/models/gr/authentication.py b/generalresearch/models/gr/authentication.py index 67a8fc2..f9644fe 100644 --- a/generalresearch/models/gr/authentication.py +++ b/generalresearch/models/gr/authentication.py @@ -3,7 +3,7 @@ from __future__ import annotations import binascii import json import os -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Self from pydantic import ( @@ -181,7 +181,7 @@ class GRUser(BaseModel): team_products = pm.fetch_uuids(team_uuids=team_uuids) if team_uuids else [] products = {p.id: p for p in business_products + team_products} - self.products = sorted(products.values(), key=lambda x: getattr(x, "created")) + self.products = sorted(products.values(), key=lambda x: x.created) def prefetch_token(self, pg_config: PostgresConfig): from generalresearch.managers.gr.authentication import ( @@ -283,7 +283,6 @@ class GRUser(BaseModel): ex=ex_secs, ) - return None # --- ORM --- diff --git a/generalresearch/models/gr/business.py b/generalresearch/models/gr/business.py index 534b23f..74b5c29 100644 --- a/generalresearch/models/gr/business.py +++ b/generalresearch/models/gr/business.py @@ -420,7 +420,7 @@ class Business(BaseModel): # that is still valid. Don't attempt to build a balance, leave it # as None rather than all zeros LOG.warning(f"Business({self.uuid=}).prebuild_balance empty dataframe") - return None + return LOG.debug(f"Business.prebuild_balance.groupby() {df.head()}") df = df.groupby("account_id").sum() @@ -442,13 +442,11 @@ class Business(BaseModel): order_by=OrderBy.DESC, ) self.prebuild_payouts_total() - return None def prebuild_payouts_total(self): assert self.payouts is not None self.payouts_total = USDCent(sum([po.amount for po in self.payouts])) self.payouts_total_str = self.payouts_total.to_usd_str() - return None def prebuild_pop_financial( self, @@ -548,7 +546,6 @@ class Business(BaseModel): except Exception as e: raise OSError(f"Parquet verification failed: {e}") - return None def prebuild_enriched_wall_parquet( self, @@ -593,7 +590,6 @@ class Business(BaseModel): except Exception as e: raise OSError(f"Parquet verification failed: {e}") - return None @classmethod def required_fields(cls) -> list[str]: diff --git a/generalresearch/models/gr/team.py b/generalresearch/models/gr/team.py index 900062f..78a9ba9 100644 --- a/generalresearch/models/gr/team.py +++ b/generalresearch/models/gr/team.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import os -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Self @@ -194,7 +194,6 @@ class Team(BaseModel): except Exception as e: raise OSError(f"Parquet verification failed: {e}") - return def prebuild_enriched_wall_parquet( self, @@ -239,7 +238,6 @@ class Team(BaseModel): except Exception as e: raise OSError(f"Parquet verification failed: {e}") - return None @classmethod def required_fields(cls) -> list[str]: @@ -325,7 +323,6 @@ class Team(BaseModel): enriched_wall=enriched_wall, ) - return # --- ORM --- @@ -344,5 +341,5 @@ class Team(BaseModel): d = {val: json.loads(res[idx]) for idx, val in enumerate(keys)} return Team.model_validate(d) - except (Exception,) as e: + except Exception: return None diff --git a/generalresearch/models/legacy/questions.py b/generalresearch/models/legacy/questions.py index 8e19e57..4651ab0 100644 --- a/generalresearch/models/legacy/questions.py +++ b/generalresearch/models/legacy/questions.py @@ -106,7 +106,6 @@ class UserQuestionAnswerIn(BaseModel): if self.question_id == user_agent_qid: val = self.answer[0] # assert val == request.user_agent.to_header(): - pass return self @@ -217,7 +216,6 @@ class UserQuestionAnswers(BaseModel): # --- Prefetch --- def prefetch_user(self, um: UserManager) -> None: - from generalresearch.models.thl.user import User res: User | None = um.get_user_if_exists( product_id=self.product_id, product_user_id=self.product_user_id @@ -230,7 +228,6 @@ class UserQuestionAnswers(BaseModel): def prefetch_wall(self, wm: WallManager) -> None: from generalresearch.models import Source - from generalresearch.models.thl.session import Wall res: Wall | None = wm.get_from_uuid_if_exists(wall_uuid=self.session_id) diff --git a/generalresearch/models/network/mtr/execute.py b/generalresearch/models/network/mtr/execute.py index 953124d..c5b3c5c 100644 --- a/generalresearch/models/network/mtr/execute.py +++ b/generalresearch/models/network/mtr/execute.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/models/network/nmap/parser.py b/generalresearch/models/network/nmap/parser.py index 6ad4ab4..ecaf2d1 100644 --- a/generalresearch/models/network/nmap/parser.py +++ b/generalresearch/models/network/nmap/parser.py @@ -1,7 +1,7 @@ from __future__ import annotations import xml.etree.ElementTree as ET -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any from generalresearch.models.network.definitions import IPProtocol diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py index 1d74df2..d6de84b 100644 --- a/generalresearch/models/network/rdns/execute.py +++ b/generalresearch/models/network/rdns/execute.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/models/spectrum/survey.py b/generalresearch/models/spectrum/survey.py index aaf182e..f9a5e27 100644 --- a/generalresearch/models/spectrum/survey.py +++ b/generalresearch/models/spectrum/survey.py @@ -75,7 +75,7 @@ class SpectrumCondition(MarketplaceCondition): rs["from"] = round(rs["from"] / 12) rs["to"] = round(rs["to"] / 12) d["values"] = [ - "{0}-{1}".format(rs["from"] or "inf", rs["to"] or "inf") + f"{rs["from"] or "inf"}-{rs["to"] or "inf"}" for rs in d["range_sets"] ] d["value_type"] = ConditionValueType.RANGE diff --git a/generalresearch/models/string_utils.py b/generalresearch/models/string_utils.py index d76456f..dff2f4d 100644 --- a/generalresearch/models/string_utils.py +++ b/generalresearch/models/string_utils.py @@ -1,5 +1,4 @@ import unicodedata -from typing import Optional def remove_nbsp(s: str | None) -> str | None: diff --git a/generalresearch/models/thl/contest/contest.py b/generalresearch/models/thl/contest/contest.py index 8c45f13..2a8853d 100644 --- a/generalresearch/models/thl/contest/contest.py +++ b/generalresearch/models/thl/contest/contest.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from abc import ABC, abstractmethod -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any, Self from uuid import uuid4 @@ -167,7 +167,7 @@ class Contest(ContestBase): ended_at=datetime.now(tz=UTC), end_reason=reason, ) - return None + return def model_dump_mysql(self, **kwargs) -> dict[str, Any]: d = self.model_dump(mode="json", **kwargs) diff --git a/generalresearch/models/thl/contest/contest_entry.py b/generalresearch/models/thl/contest/contest_entry.py index 31ef317..bb3aef4 100644 --- a/generalresearch/models/thl/contest/contest_entry.py +++ b/generalresearch/models/thl/contest/contest_entry.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 from pydantic import ( diff --git a/generalresearch/models/thl/contest/examples.py b/generalresearch/models/thl/contest/examples.py index 1810e63..018f090 100644 --- a/generalresearch/models/thl/contest/examples.py +++ b/generalresearch/models/thl/contest/examples.py @@ -93,7 +93,6 @@ def _example_raffle(schema: dict) -> None: product_id=EXAMPLE_PRODUCT_ID, ).model_dump(mode="json") - return None def _example_raffle_user_view(schema: dict[str, Any]) -> None: @@ -144,7 +143,6 @@ def _example_raffle_user_view(schema: dict[str, Any]) -> None: product_user_id="test-user", ).model_dump(mode="json") - return None def _example_milestone_create(schema: dict[str, Any]) -> None: @@ -185,7 +183,6 @@ def _example_milestone_create(schema: dict[str, Any]) -> None: terms_and_conditions=HttpUrl("https://www.example.com"), ).model_dump(mode="json") - return None def _example_milestone(schema: dict[str, Any]) -> None: @@ -231,7 +228,6 @@ def _example_milestone(schema: dict[str, Any]) -> None: win_count=12, ).model_dump(mode="json") - return None def _example_milestone_user_view(schema: dict[str, Any]) -> None: @@ -277,7 +273,6 @@ def _example_milestone_user_view(schema: dict[str, Any]) -> None: product_user_id="test-user", ).model_dump(mode="json") - return None def _example_leaderboard_contest_create(schema: dict[str, Any]) -> None: @@ -322,7 +317,6 @@ def _example_leaderboard_contest_create(schema: dict[str, Any]) -> None: leaderboard_key=f"leaderboard:{EXAMPLE_PRODUCT_ID}:us:weekly:2025-05-26:complete_count", ).model_dump(mode="json") - return None def _example_leaderboard_contest(schema: dict[str, Any]) -> None: @@ -368,7 +362,6 @@ def _example_leaderboard_contest(schema: dict[str, Any]) -> None: product_id=EXAMPLE_PRODUCT_ID, ).model_dump(mode="json") - return None def _example_leaderboard_contest_user_view(schema: dict[str, Any]) -> None: diff --git a/generalresearch/models/thl/contest/io.py b/generalresearch/models/thl/contest/io.py index c6af719..d11080f 100644 --- a/generalresearch/models/thl/contest/io.py +++ b/generalresearch/models/thl/contest/io.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 from generalresearch.models.thl.contest.definitions import ContestType diff --git a/generalresearch/models/thl/contest/leaderboard.py b/generalresearch/models/thl/contest/leaderboard.py index c5e0626..696cdea 100644 --- a/generalresearch/models/thl/contest/leaderboard.py +++ b/generalresearch/models/thl/contest/leaderboard.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from typing import Any, Literal, Self from pydantic import ( diff --git a/generalresearch/models/thl/contest/milestone.py b/generalresearch/models/thl/contest/milestone.py index 8b74d50..8d96fcb 100644 --- a/generalresearch/models/thl/contest/milestone.py +++ b/generalresearch/models/thl/contest/milestone.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from datetime import timedelta -from typing import Any, Literal +from typing import Any, Literal, Self from pydantic import ( BaseModel, @@ -10,7 +10,6 @@ from pydantic import ( Field, PositiveInt, ) -from typing import Self from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.thl.contest.contest import ( diff --git a/generalresearch/models/thl/contest/raffle.py b/generalresearch/models/thl/contest/raffle.py index b497a44..08243f4 100644 --- a/generalresearch/models/thl/contest/raffle.py +++ b/generalresearch/models/thl/contest/raffle.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging import random from collections import defaultdict -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any, Literal, Self from pydantic import ( diff --git a/generalresearch/models/thl/finance.py b/generalresearch/models/thl/finance.py index a992f78..0856825 100644 --- a/generalresearch/models/thl/finance.py +++ b/generalresearch/models/thl/finance.py @@ -1,7 +1,7 @@ from __future__ import annotations import random -from datetime import UTC, timezone +from datetime import UTC from typing import TYPE_CHECKING from uuid import uuid4 @@ -28,8 +28,7 @@ payout_example = random.randint(150, 750 * 100) adjustment_example = random.randint(-1_000, 50 * 100) if TYPE_CHECKING: - from generalresearch.managers.thl.product import ProductManager - from generalresearch.models.thl.ledger import AccountType, Direction, LedgerAccount + from generalresearch.models.thl.ledger import LedgerAccount class AdjustmentType(BaseModel): diff --git a/generalresearch/models/thl/ipinfo.py b/generalresearch/models/thl/ipinfo.py index 3f212cf..e327bae 100644 --- a/generalresearch/models/thl/ipinfo.py +++ b/generalresearch/models/thl/ipinfo.py @@ -1,7 +1,7 @@ from __future__ import annotations import ipaddress -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any, Literal, Self from faker import Faker diff --git a/generalresearch/models/thl/ledger_example.py b/generalresearch/models/thl/ledger_example.py index 92ad83d..0291691 100644 --- a/generalresearch/models/thl/ledger_example.py +++ b/generalresearch/models/thl/ledger_example.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any from uuid import uuid4 diff --git a/generalresearch/models/thl/offerwall/cache.py b/generalresearch/models/thl/offerwall/cache.py index 82ab36d..97546b2 100644 --- a/generalresearch/models/thl/offerwall/cache.py +++ b/generalresearch/models/thl/offerwall/cache.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any from pydantic import BaseModel, Field diff --git a/generalresearch/models/thl/payout_format.py b/generalresearch/models/thl/payout_format.py index a7fc4fb..4d616b6 100644 --- a/generalresearch/models/thl/payout_format.py +++ b/generalresearch/models/thl/payout_format.py @@ -45,7 +45,7 @@ def format_payout_format(payout_format: str, payout_int: int) -> str: try: xform, formatstr = inside.split(":") - except ValueError as e: + except ValueError: raise ValueError( "Payout format string must contain ':' to distinguish between transformations and formatting." ) @@ -61,17 +61,17 @@ def format_payout_format(payout_format: str, payout_int: int) -> str: payout = decimal.Decimal(eval(xform, {"payout": payout_int})) - except NameError as e: + except NameError: raise ValueError("Payout format string must contain 'payout' variable.") - except ZeroDivisionError as e: + except ZeroDivisionError: raise ValueError("Cannot divide by zero.") - except TypeError as e: + except TypeError: # "{payout()*1:}" - TypeError: 'int' object is not callable raise ValueError("Invalid type reference.") - except Exception as e: - raise ValueError(f"Invalid payout transformation") + except Exception: + raise ValueError("Invalid payout transformation") formatstr = f"{{:{formatstr}}}" diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index 5b7ab9a..76a8e83 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -77,7 +77,6 @@ if TYPE_CHECKING: from generalresearch.models.thl.payout import ( BrokerageProductPayoutEvent, ) - from generalresearch.models.thl.user import User # fmt: off @@ -1088,7 +1087,6 @@ class Product(BaseModel, validate_assignment=True): from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) - from generalresearch.models.thl.ledger import LedgerAccount account: LedgerAccount = thl_lm.get_account_or_create_bp_wallet(product=self) assert self.id == account.reference_uuid diff --git a/generalresearch/models/thl/profiling/marketplace.py b/generalresearch/models/thl/profiling/marketplace.py index 9038cf6..ad4ce80 100644 --- a/generalresearch/models/thl/profiling/marketplace.py +++ b/generalresearch/models/thl/profiling/marketplace.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from functools import cached_property from typing import Any diff --git a/generalresearch/models/thl/profiling/upk_question.py b/generalresearch/models/thl/profiling/upk_question.py index 08ea350..77bba6f 100644 --- a/generalresearch/models/thl/profiling/upk_question.py +++ b/generalresearch/models/thl/profiling/upk_question.py @@ -367,7 +367,7 @@ class UpkQuestion(BaseModel): self.choices is None ), f"No `choices` are allowed for type `{self.type}`" else: - assert self.choices is not None, f"`choices` must be set" + assert self.choices is not None, "`choices` must be set" return self @model_validator(mode="after") diff --git a/generalresearch/models/thl/profiling/upk_question_answer.py b/generalresearch/models/thl/profiling/upk_question_answer.py index c59d99d..d8323ad 100644 --- a/generalresearch/models/thl/profiling/upk_question_answer.py +++ b/generalresearch/models/thl/profiling/upk_question_answer.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any, Self from uuid import uuid4 diff --git a/generalresearch/models/thl/profiling/user_question_answer.py b/generalresearch/models/thl/profiling/user_question_answer.py index a55b205..2db07b7 100644 --- a/generalresearch/models/thl/profiling/user_question_answer.py +++ b/generalresearch/models/thl/profiling/user_question_answer.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from collections.abc import Iterator -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from typing import Any, Literal, Self from pydantic import ( diff --git a/generalresearch/models/thl/session.py b/generalresearch/models/thl/session.py index c8e681c..fe7194a 100644 --- a/generalresearch/models/thl/session.py +++ b/generalresearch/models/thl/session.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import logging -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import TYPE_CHECKING, Annotated, Any, Self from uuid import uuid4 diff --git a/generalresearch/models/thl/survey/__init__.py b/generalresearch/models/thl/survey/__init__.py index bc3d699..6e6b475 100644 --- a/generalresearch/models/thl/survey/__init__.py +++ b/generalresearch/models/thl/survey/__init__.py @@ -111,7 +111,6 @@ class MarketplaceTask(BaseModel, ABC): """ The Condition Model for this survey class """ - pass @property @abstractmethod @@ -119,7 +118,6 @@ class MarketplaceTask(BaseModel, ABC): """ The age question ID """ - pass @property @abstractmethod @@ -129,7 +127,6 @@ class MarketplaceTask(BaseModel, ABC): """ Mapping of generic Gender to the marketplace condition for that gender """ - pass @property def marketplace_age_groups( diff --git a/generalresearch/models/thl/survey/buyer.py b/generalresearch/models/thl/survey/buyer.py index 384bab4..b888007 100644 --- a/generalresearch/models/thl/survey/buyer.py +++ b/generalresearch/models/thl/survey/buyer.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from math import log from typing import Annotated diff --git a/generalresearch/models/thl/survey/model.py b/generalresearch/models/thl/survey/model.py index 57bcbe2..2eed8f7 100644 --- a/generalresearch/models/thl/survey/model.py +++ b/generalresearch/models/thl/survey/model.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from typing import Annotated, Any diff --git a/generalresearch/models/thl/survey/penalty.py b/generalresearch/models/thl/survey/penalty.py index 05153fe..04f8e20 100644 --- a/generalresearch/models/thl/survey/penalty.py +++ b/generalresearch/models/thl/survey/penalty.py @@ -1,7 +1,7 @@ from __future__ import annotations import abc -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, TypeAdapter diff --git a/generalresearch/models/thl/task_adjustment.py b/generalresearch/models/thl/task_adjustment.py index 1834898..27c47d4 100644 --- a/generalresearch/models/thl/task_adjustment.py +++ b/generalresearch/models/thl/task_adjustment.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from uuid import uuid4 diff --git a/generalresearch/models/thl/user.py b/generalresearch/models/thl/user.py index 55bbd18..11e0d67 100644 --- a/generalresearch/models/thl/user.py +++ b/generalresearch/models/thl/user.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import logging import re -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import TYPE_CHECKING, Annotated, Self from uuid import UUID, uuid4 diff --git a/generalresearch/models/thl/user_iphistory.py b/generalresearch/models/thl/user_iphistory.py index 469f8ba..257c41a 100644 --- a/generalresearch/models/thl/user_iphistory.py +++ b/generalresearch/models/thl/user_iphistory.py @@ -1,7 +1,7 @@ from __future__ import annotations import ipaddress -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from typing import Self from faker import Faker @@ -204,7 +204,6 @@ class UserIPHistory(BaseModel): if res.get(x.ip): x.information = res[x.ip] - return None def collapse_ip_records(self): """ diff --git a/generalresearch/models/thl/wallet/payout.py b/generalresearch/models/thl/wallet/payout.py index cbb37fe..42530b3 100644 --- a/generalresearch/models/thl/wallet/payout.py +++ b/generalresearch/models/thl/wallet/payout.py @@ -2,7 +2,7 @@ from __future__ import annotations import json from collections.abc import Collection -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import Any from uuid import uuid4 diff --git a/generalresearch/pg_helper.py b/generalresearch/pg_helper.py index b5e124a..1d5d30b 100644 --- a/generalresearch/pg_helper.py +++ b/generalresearch/pg_helper.py @@ -1,12 +1,11 @@ from __future__ import annotations -from datetime import UTC, timezone +from datetime import UTC import psycopg -from psycopg.adapt import Buffer from psycopg.rows import RowFactory, dict_row from psycopg.types.datetime import TimestampLoader -from psycopg.types.net import Address, InetLoader, Interface +from psycopg.types.net import InetLoader from psycopg.types.string import TextLoader from psycopg.types.uuid import UUIDLoader from pydantic import PostgresDsn @@ -103,10 +102,9 @@ class PostgresConfig: # This is only intended for SELECT queries assert "SELECT" in query.upper(), "Supports SELECTs only" - with self.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=query, params=params) - return c.fetchall() + with self.make_connection() as conn, conn.cursor() as c: + c.execute(query=query, params=params) + return c.fetchall() def execute_write(self, query, params=None) -> int: cmd = query.lstrip().upper() diff --git a/generalresearch/schemas/survey_stats.py b/generalresearch/schemas/survey_stats.py index 0d509e1..b3acf34 100644 --- a/generalresearch/schemas/survey_stats.py +++ b/generalresearch/schemas/survey_stats.py @@ -70,7 +70,7 @@ UnitInterval = Column( SID_CHECKS = [ Check.str_length(min_value=3, max_value=67), - Check.str_matches("^[a-z]{1,2}\:[A-Za-z0-9]+"), + Check.str_matches(r"^[a-z]{1,2}\:[A-Za-z0-9]+"), Check( lambda x: len(set(x.str.split(":").str[0])) == 1, error="the sources must all be the same", diff --git a/generalresearch/sql_helper.py b/generalresearch/sql_helper.py index 2bbdbe5..ea45305 100644 --- a/generalresearch/sql_helper.py +++ b/generalresearch/sql_helper.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import Any, Optional +from typing import Any from uuid import UUID from pydantic import MariaDBDsn, MySQLDsn, PostgresDsn @@ -39,14 +39,7 @@ class SqlConnector: # I'm intentionally doing a match case here so that we'll make sure # we can NOT use this on old versions of python 😈 - if "mysql" in self.dsn.scheme: - import pymysql as engine_module - - self.engine_module = engine_module - self.cursor_class = engine_module.cursors.DictCursor - self.quote_char = "`" - - elif "maria" in self.dsn.scheme: + if "mysql" in self.dsn.scheme or "maria" in self.dsn.scheme: import pymysql as engine_module self.engine_module = engine_module @@ -199,7 +192,6 @@ class SqlHelper(SqlConnector): if cursor is None: c.connection.commit() - return None def bulk_update( self, @@ -233,7 +225,7 @@ class SqlHelper(SqlConnector): if cursor is None: c.connection.commit() - return None + return def get_or_create( self, @@ -253,7 +245,7 @@ class SqlHelper(SqlConnector): lookup_fns = ",".join( ["`" + x + "`" for x in set(lookup_dict.keys()) | {primary_key}] ) - lookup_vals = " AND ".join([f"`{fn}`=%({fn})s" for fn in lookup_dict.keys()]) + lookup_vals = " AND ".join([f"`{fn}`=%({fn})s" for fn in lookup_dict]) table_name_str = self._quote(table_name) query = f"SELECT {lookup_fns} FROM {table_name_str} WHERE {lookup_vals} LIMIT 2" if cursor is None: @@ -293,7 +285,7 @@ class SqlHelper(SqlConnector): else: c = cursor field_names = ",".join(map(self._quote, create_dict)) - vals = ",".join([f"%({fn})s" for fn in create_dict.keys()]) + vals = ",".join([f"%({fn})s" for fn in create_dict]) table_name_str = self._quote(table_name) query = f"INSERT INTO {table_name_str} ({field_names}) VALUES ({vals})" c.execute(query, create_dict) @@ -352,4 +344,3 @@ class SqlHelper(SqlConnector): if cursor is None: c.connection.commit() - return None diff --git a/generalresearch/thl_django/apps.py b/generalresearch/thl_django/apps.py index 2813947..bd87110 100644 --- a/generalresearch/thl_django/apps.py +++ b/generalresearch/thl_django/apps.py @@ -6,11 +6,11 @@ class THLSchemaConfig(AppConfig): label = "thl_django" def ready(self): - from .accounting import models # noqa: F401 # pycharm: keep - from .common import models # noqa: F401 # pycharm: keep - from .contest import models # noqa: F401 # pycharm: keep - from .event import models # noqa: F401 # pycharm: keep - from .marketplace import models # noqa: F401 # pycharm: keep - from .network import models # noqa: F401 # pycharm: keep - from .userhealth import models # noqa: F401 # pycharm: keep + from .accounting import models # pycharm: keep + from .common import models # pycharm: keep + from .contest import models # pycharm: keep + from .event import models # pycharm: keep + from .marketplace import models # pycharm: keep + from .network import models # pycharm: keep + from .userhealth import models # pycharm: keep from .userprofile import models # noqa: F401 # pycharm: keep diff --git a/generalresearch/thl_django/fields.py b/generalresearch/thl_django/fields.py index 5e40ef0..251faa5 100644 --- a/generalresearch/thl_django/fields.py +++ b/generalresearch/thl_django/fields.py @@ -1,6 +1,7 @@ -from django.db import models import ipaddress +from django.db import models + class CIDRField(models.Field): description = "PostgreSQL CIDR network" diff --git a/generalresearch/thl_django/migrations/0001_initial.py b/generalresearch/thl_django/migrations/0001_initial.py index ecae35a..cf147e3 100644 --- a/generalresearch/thl_django/migrations/0001_initial.py +++ b/generalresearch/thl_django/migrations/0001_initial.py @@ -1,7 +1,8 @@ # Generated by Django 6.0 on 2025-12-26 20:53 -import django.db.models.deletion import uuid + +import django.db.models.deletion from django.db import migrations, models diff --git a/generalresearch/thl_django/migrations/0002_surveystat_is_live_alter_surveycategory_strength_and_more.py b/generalresearch/thl_django/migrations/0002_surveystat_is_live_alter_surveycategory_strength_and_more.py index 211c48a..f767afc 100644 --- a/generalresearch/thl_django/migrations/0002_surveystat_is_live_alter_surveycategory_strength_and_more.py +++ b/generalresearch/thl_django/migrations/0002_surveystat_is_live_alter_surveycategory_strength_and_more.py @@ -1,7 +1,7 @@ # Generated by Django 6.0 on 2025-12-28 16:49 -from django.db import migrations, models from django.contrib.postgres.operations import AddIndexConcurrently +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/generalresearch/thl_django/migrations/0003_remove_surveystat_surveystat_live_survey_idx_and_more.py b/generalresearch/thl_django/migrations/0003_remove_surveystat_surveystat_live_survey_idx_and_more.py index ecaf0a9..dcf9ef2 100644 --- a/generalresearch/thl_django/migrations/0003_remove_surveystat_surveystat_live_survey_idx_and_more.py +++ b/generalresearch/thl_django/migrations/0003_remove_surveystat_surveystat_live_survey_idx_and_more.py @@ -1,10 +1,10 @@ # Generated by Django 6.0 on 2025-12-29 21:22 -from django.db import migrations, models from django.contrib.postgres.operations import ( AddIndexConcurrently, RemoveIndexConcurrently, ) +from django.db import migrations, models class Migration(migrations.Migration): diff --git a/generalresearch/thl_django/migrations/0006_remove_thlsession_thl_session_status_d578b7_idx_and_more.py b/generalresearch/thl_django/migrations/0006_remove_thlsession_thl_session_status_d578b7_idx_and_more.py index e2492ab..64338c8 100644 --- a/generalresearch/thl_django/migrations/0006_remove_thlsession_thl_session_status_d578b7_idx_and_more.py +++ b/generalresearch/thl_django/migrations/0006_remove_thlsession_thl_session_status_d578b7_idx_and_more.py @@ -1,7 +1,7 @@ # Generated by Django 6.0 on 2026-01-02 17:38 -from django.db import migrations from django.contrib.postgres.operations import RemoveIndexConcurrently +from django.db import migrations class Migration(migrations.Migration): diff --git a/generalresearch/thl_django/migrations/0009_toolrun_mtrhop_portscanport_iplabel_mtr_portscan_and_more.py b/generalresearch/thl_django/migrations/0009_toolrun_mtrhop_portscanport_iplabel_mtr_portscan_and_more.py index e19a353..e8ac2c2 100644 --- a/generalresearch/thl_django/migrations/0009_toolrun_mtrhop_portscanport_iplabel_mtr_portscan_and_more.py +++ b/generalresearch/thl_django/migrations/0009_toolrun_mtrhop_portscanport_iplabel_mtr_portscan_and_more.py @@ -1,13 +1,14 @@ # Generated by Django 6.0 on 2026-03-15 20:17 +import uuid + import django.contrib.postgres.indexes import django.db.models.deletion import django.utils.timezone from django.contrib.postgres.operations import CreateExtension +from django.db import migrations, models import generalresearch.thl_django.fields -import uuid -from django.db import migrations, models class Migration(migrations.Migration): diff --git a/generalresearch/thl_django/network/models.py b/generalresearch/thl_django/network/models.py index 167af02..733c0ab 100644 --- a/generalresearch/thl_django/network/models.py +++ b/generalresearch/thl_django/network/models.py @@ -1,12 +1,11 @@ from uuid import uuid4 -from django.utils import timezone -from django.contrib.postgres.indexes import GistIndex, GinIndex +from django.contrib.postgres.indexes import GinIndex, GistIndex from django.db import models +from django.utils import timezone from generalresearch.thl_django.fields import CIDRField - ####### # ** Signals ** # ToolRun diff --git a/generalresearch/utils/enum.py b/generalresearch/utils/enum.py index e59a383..b4620b3 100644 --- a/generalresearch/utils/enum.py +++ b/generalresearch/utils/enum.py @@ -19,7 +19,7 @@ class ReprEnumMeta(EnumMeta): [f" - __{e.value}__ *({e.name})*: {descriptions[e.name]}" for e in self] ) else: - return f"\nAllowed values: \n" + "\n".join( + return "\nAllowed values: \n" + "\n".join( [f" - __{e.value}__ *({e.name})*: {descriptions[e.name]}" for e in self] ) @@ -35,7 +35,7 @@ class ReprEnumMeta(EnumMeta): [f" - __{e.name}__: {descriptions[e.name]}" for e in self] ) else: - return f"\nAllowed values: \n" + "\n".join( + return "\nAllowed values: \n" + "\n".join( [f" - __{e.name}__: {descriptions[e.name]}" for e in self] ) diff --git a/generalresearch/wall_status_codes/lucid.py b/generalresearch/wall_status_codes/lucid.py index 05b75fb..3cc1b5e 100644 --- a/generalresearch/wall_status_codes/lucid.py +++ b/generalresearch/wall_status_codes/lucid.py @@ -61,7 +61,7 @@ client_status_map: dict[str, StatusCode1] = { "35": StatusCode1.BUYER_QUALITY_FAIL, } -status_map = defaultdict(lambda: Status.FAIL, **{"s": Status.COMPLETE}) +status_map = defaultdict(lambda: Status.FAIL, s=Status.COMPLETE) status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: [], StatusCode1.BUYER_FAIL: ["3"], diff --git a/generalresearch/wall_status_codes/morning.py b/generalresearch/wall_status_codes/morning.py index bac9669..ffa6be2 100644 --- a/generalresearch/wall_status_codes/morning.py +++ b/generalresearch/wall_status_codes/morning.py @@ -52,7 +52,7 @@ short_code_to_status_codes_morning: dict[str, str] = { "sur_tim": "survey_timeout", "tem_ban": "temporarily_banned", } -status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) +status_map = defaultdict(lambda: Status.FAIL, complete=Status.COMPLETE) status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["complete"], diff --git a/generalresearch/wall_status_codes/pollfish.py b/generalresearch/wall_status_codes/pollfish.py index 1ae732d..a5c6e25 100644 --- a/generalresearch/wall_status_codes/pollfish.py +++ b/generalresearch/wall_status_codes/pollfish.py @@ -28,7 +28,7 @@ status_codes_map: dict[str, str] = { "su_al_ta": "survey_already_taken", "complete": "complete", } -status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) +status_map = defaultdict(lambda: Status.FAIL, complete=Status.COMPLETE) status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["complete"], StatusCode1.BUYER_FAIL: ["third_party_termination", "screenout"], diff --git a/generalresearch/wall_status_codes/precision.py b/generalresearch/wall_status_codes/precision.py index b033cf2..c7593cf 100644 --- a/generalresearch/wall_status_codes/precision.py +++ b/generalresearch/wall_status_codes/precision.py @@ -45,7 +45,7 @@ status_codes_precision: dict[str, str] = { "60": "Client Reject", "80": "Final Complete", } -status_map = defaultdict(lambda: Status.FAIL, **{"s": Status.COMPLETE}) +status_map = defaultdict(lambda: Status.FAIL, s=Status.COMPLETE) status_codes_ext_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["10"], StatusCode1.BUYER_FAIL: ["20", "30"], diff --git a/generalresearch/wall_status_codes/repdata.py b/generalresearch/wall_status_codes/repdata.py index ad64338..24aab14 100644 --- a/generalresearch/wall_status_codes/repdata.py +++ b/generalresearch/wall_status_codes/repdata.py @@ -46,7 +46,7 @@ rd_threat_name: dict[str, str] = { "18": "MaxMind Failure", } -status_map = defaultdict(lambda: Status.FAIL, **{"complete": Status.COMPLETE}) +status_map = defaultdict(lambda: Status.FAIL, complete=Status.COMPLETE) status_code_map: dict[StatusCode1, list[str]] = { StatusCode1.COMPLETE: ["1000"], StatusCode1.BUYER_FAIL: ["2000", "4000"], diff --git a/test_utils/conftest.py b/test_utils/conftest.py index ffe458c..03e9305 100644 --- a/test_utils/conftest.py +++ b/test_utils/conftest.py @@ -5,9 +5,8 @@ import shutil import stat import subprocess import sys -import tempfile from collections.abc import Callable, Generator -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from os.path import join as pjoin from pathlib import Path from uuid import uuid4 diff --git a/test_utils/grliq/conftest.py b/test_utils/grliq/conftest.py index 7665b52..9a3bc56 100644 --- a/test_utils/grliq/conftest.py +++ b/test_utils/grliq/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import uuid4 import pytest diff --git a/test_utils/incite/conftest.py b/test_utils/incite/conftest.py index 87ea7ae..2968d18 100644 --- a/test_utils/incite/conftest.py +++ b/test_utils/incite/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from os.path import join as pjoin from pathlib import Path from random import choice as randchoice diff --git a/test_utils/managers/conftest.py b/test_utils/managers/conftest.py index 4dacb29..2a9ea00 100644 --- a/test_utils/managers/conftest.py +++ b/test_utils/managers/conftest.py @@ -20,14 +20,12 @@ from generalresearch.managers.thl.ipinfo import ( IPGeonameManager, IPInformationManager, ) -from generalresearch.managers.thl.profiling.uqa import UQAManager from generalresearch.managers.thl.userhealth import ( AuditLogManager, IPRecordManager, UserIpHistoryManager, ) from generalresearch.models import Source -from generalresearch.models.thl.user import User from generalresearch.pg_helper import PostgresConfig from generalresearch.redis_helper import RedisConfig from generalresearch.sql_helper import SqlHelper diff --git a/test_utils/models/conftest.py b/test_utils/models/conftest.py index 3a9e45c..5570b40 100644 --- a/test_utils/models/conftest.py +++ b/test_utils/models/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from random import choice as randchoice from random import randint @@ -19,17 +19,15 @@ from generalresearch.models.thl.definitions import ( ) from generalresearch.models.thl.survey.model import Buyer, Survey from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig if TYPE_CHECKING: from generalresearch.currency import USDCent - from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager from generalresearch.managers.gr.business import ( BusinessAddressManager, BusinessBankAccountManager, BusinessManager, ) - from generalresearch.managers.gr.team import MembershipManager, TeamManager + from generalresearch.managers.gr.team import TeamManager from generalresearch.managers.thl.buyer import BuyerManager from generalresearch.managers.thl.ipinfo import ( IPGeonameManager, @@ -45,13 +43,12 @@ if TYPE_CHECKING: from generalresearch.managers.thl.user_manager.user_manager import UserManager from generalresearch.managers.thl.userhealth import AuditLogManager, IPRecordManager from generalresearch.managers.thl.wall import WallManager - from generalresearch.models.gr.authentication import GRToken, GRUser from generalresearch.models.gr.business import ( Business, BusinessAddress, BusinessBankAccount, ) - from generalresearch.models.gr.team import Membership, Team + from generalresearch.models.gr.team import Team from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation from generalresearch.models.thl.payout import ( BrokerageProductPayoutEvent, diff --git a/test_utils/models/contest/conftest.py b/test_utils/models/contest/conftest.py index 0060946..b14e126 100644 --- a/test_utils/models/contest/conftest.py +++ b/test_utils/models/contest/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from uuid import uuid4 diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py index cabd8dc..6ba37a3 100644 --- a/test_utils/models/network/conftest.py +++ b/test_utils/models/network/conftest.py @@ -1,5 +1,5 @@ import os -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from uuid import uuid4 import pytest diff --git a/test_utils/models/thl/conftest.py b/test_utils/models/thl/conftest.py index a2adcce..5b21c6b 100644 --- a/test_utils/models/thl/conftest.py +++ b/test_utils/models/thl/conftest.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import ROUND_DOWN, Decimal from random import choice as rand_choice from random import choice as rchoice diff --git a/test_utils/spectrum/conftest.py b/test_utils/spectrum/conftest.py index 9c067d3..eb2e289 100644 --- a/test_utils/spectrum/conftest.py +++ b/test_utils/spectrum/conftest.py @@ -1,6 +1,6 @@ import logging import time -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import TYPE_CHECKING import pytest diff --git a/tests/incite/collections/test_df_collection_base.py b/tests/incite/collections/test_df_collection_base.py index c3c64e4..176bf4b 100644 --- a/tests/incite/collections/test_df_collection_base.py +++ b/tests/incite/collections/test_df_collection_base.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import TYPE_CHECKING import pandas as pd @@ -9,7 +9,6 @@ from generalresearch.incite.collections import ( DFCollection, DFCollectionType, ) -from test_utils.incite.conftest import mnt_filepath if TYPE_CHECKING: from generalresearch.incite.base import GRLDatasets diff --git a/tests/incite/collections/test_df_collection_item_base.py b/tests/incite/collections/test_df_collection_item_base.py index 8cf719d..0218f30 100644 --- a/tests/incite/collections/test_df_collection_item_base.py +++ b/tests/incite/collections/test_df_collection_item_base.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from typing import TYPE_CHECKING import pytest diff --git a/tests/incite/collections/test_df_collection_item_thl_web.py b/tests/incite/collections/test_df_collection_item_thl_web.py index 062171d..3d70e56 100644 --- a/tests/incite/collections/test_df_collection_item_thl_web.py +++ b/tests/incite/collections/test_df_collection_item_thl_web.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Callable, Generator -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product as iter_product from os.path import join as pjoin from pathlib import Path, PurePath @@ -236,7 +236,6 @@ class TestDFCollectionItemMethod: incite_item_factory, delete_df_collection, ): - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -280,7 +279,6 @@ class TestDFCollectionItemMethod: incite_item_factory, delete_df_collection, ): - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -333,7 +331,6 @@ class TestDFCollectionItemMethod: delete_df_collection, mnt_filepath, ): - from generalresearch.models.thl.user import User if df_collection.data_type != DFCollectionType.LEDGER: return @@ -385,7 +382,6 @@ class TestDFCollectionItemMethod: delete_df_collection, mnt_filepath, ): - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -766,7 +762,6 @@ class TestDFCollectionItemFunctionalTest: delete_df_collection, mnt_filepath: GRLDatasets, ): - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -828,7 +823,6 @@ class TestDFCollectionItemFunctionalTest: import pyarrow.parquet as pq - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -864,7 +858,6 @@ class TestDFCollectionItemFunctionalTest: delete_df_collection, mnt_filepath: GRLDatasets, ): - from generalresearch.models.thl.user import User delete_df_collection(coll=df_collection) df_collection._client = client_no_amm @@ -920,7 +913,6 @@ class TestDFCollectionItemFunctionalTest: """Don't allow creating an archive for data that will likely be overwritten or updated """ - from generalresearch.models.thl.user import User if df_collection.data_type in unsupported_mock_types: return @@ -987,7 +979,6 @@ class TestDFCollectionItemFunctionalTest: duration: timedelta, mnt_filepath, ): - from generalresearch.models.thl.user import User delete_df_collection(coll=df_collection) df_collection._client = client_no_amm diff --git a/tests/incite/collections/test_df_collection_thl_marketplaces.py b/tests/incite/collections/test_df_collection_thl_marketplaces.py index 2597d38..d2d3ce4 100644 --- a/tests/incite/collections/test_df_collection_thl_marketplaces.py +++ b/tests/incite/collections/test_df_collection_thl_marketplaces.py @@ -1,6 +1,5 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from itertools import product -from typing import TYPE_CHECKING import pytest from pandera.pandas import Column, DataFrameSchema, Index @@ -12,10 +11,6 @@ from generalresearch.incite.collections.thl_marketplaces import ( SagoSurveyHistoryCollection, SpectrumSurveyTimeseriesCollection, ) -from test_utils.incite.conftest import mnt_filepath - -if TYPE_CHECKING: - from generalresearch.incite.base import GRLDatasets def combo_object(): diff --git a/tests/incite/collections/test_df_collection_thl_web.py b/tests/incite/collections/test_df_collection_thl_web.py index 2cb0ba0..bcdeb83 100644 --- a/tests/incite/collections/test_df_collection_thl_web.py +++ b/tests/incite/collections/test_df_collection_thl_web.py @@ -13,9 +13,7 @@ from pandera.pandas import DataFrameSchema from generalresearch.incite.collections import DFCollection, DFCollectionType if TYPE_CHECKING: - from generalresearch.incite.base import GRLDatasets from generalresearch.incite.collections import ( - DFCollectionItem, DFCollectionType, ) diff --git a/tests/incite/mergers/foundations/test_enriched_session.py b/tests/incite/mergers/foundations/test_enriched_session.py index a0ae01e..ba11725 100644 --- a/tests/incite/mergers/foundations/test_enriched_session.py +++ b/tests/incite/mergers/foundations/test_enriched_session.py @@ -1,7 +1,6 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from itertools import product -from typing import Optional import dask.dataframe as dd import pandas as pd @@ -11,10 +10,6 @@ from generalresearch.incite.schemas.admin_responses import ( AdminPOPSessionSchema, ) from generalresearch.pg_helper import PostgresConfig -from test_utils.incite.collections.conftest import ( - session_collection, - wall_collection, -) @pytest.mark.parametrize( diff --git a/tests/incite/mergers/foundations/test_enriched_task_adjust.py b/tests/incite/mergers/foundations/test_enriched_task_adjust.py index 96c214f..8c3a647 100644 --- a/tests/incite/mergers/foundations/test_enriched_task_adjust.py +++ b/tests/incite/mergers/foundations/test_enriched_task_adjust.py @@ -5,13 +5,6 @@ import dask.dataframe as dd import pandas as pd import pytest -from test_utils.incite.collections.conftest import ( - wall_collection, - task_adj_collection, - session_collection, -) -from test_utils.incite.mergers.conftest import enriched_wall_merge - @pytest.mark.parametrize( argnames="offset, duration,", diff --git a/tests/incite/mergers/foundations/test_enriched_wall.py b/tests/incite/mergers/foundations/test_enriched_wall.py index b421df8..0e28bce 100644 --- a/tests/incite/mergers/foundations/test_enriched_wall.py +++ b/tests/incite/mergers/foundations/test_enriched_wall.py @@ -1,34 +1,15 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from itertools import product as iter_product -from typing import Optional import dask.dataframe as dd import pandas as pd import pytest # noinspection PyUnresolvedReferences -from distributed.utils_test import ( - cleanup, - client, - client_no_amm, - cluster_fixture, - gen_cluster, - loop, - loop_in_thread, -) - from generalresearch.incite.mergers.foundations.enriched_wall import ( EnrichedWallMergeItem, ) -from test_utils.incite.collections.conftest import ( - session_collection, - wall_collection, -) -from test_utils.incite.conftest import incite_item_factory -from test_utils.incite.mergers.conftest import ( - enriched_wall_merge, -) @pytest.mark.parametrize( @@ -118,7 +99,7 @@ class TestEnrichedWall: try: modified_time1 = path.stat().st_mtime - except (Exception,): + except Exception: modified_time1 = 0 item.build( diff --git a/tests/incite/mergers/foundations/test_user_id_product.py b/tests/incite/mergers/foundations/test_user_id_product.py index a696b45..10802e5 100644 --- a/tests/incite/mergers/foundations/test_user_id_product.py +++ b/tests/incite/mergers/foundations/test_user_id_product.py @@ -1,24 +1,13 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product import pandas as pd import pytest # noinspection PyUnresolvedReferences -from distributed.utils_test import ( - cleanup, - client, - client_no_amm, - cluster_fixture, - gen_cluster, - loop, - loop_in_thread, -) - from generalresearch.incite.mergers.foundations.user_id_product import ( UserIdProductMergeItem, ) -from test_utils.incite.mergers.conftest import user_id_product_merge @pytest.mark.parametrize( @@ -51,7 +40,7 @@ class TestUserIDProduct: try: modified_time1 = path.stat().st_mtime - except (Exception,): + except Exception: modified_time1 = 0 user_id_product_merge.build(client=client_no_amm, user_coll=user_collection) diff --git a/tests/incite/mergers/test_merge_collection.py b/tests/incite/mergers/test_merge_collection.py index 77fa8c7..15fa4db 100644 --- a/tests/incite/mergers/test_merge_collection.py +++ b/tests/incite/mergers/test_merge_collection.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product import pandas as pd @@ -9,7 +9,6 @@ from generalresearch.incite.mergers import ( MergeCollection, MergeType, ) -from test_utils.incite.conftest import mnt_filepath merge_types = list(e for e in MergeType if e != MergeType.TEST) diff --git a/tests/incite/mergers/test_merge_collection_item.py b/tests/incite/mergers/test_merge_collection_item.py index 96f8789..3d0b644 100644 --- a/tests/incite/mergers/test_merge_collection_item.py +++ b/tests/incite/mergers/test_merge_collection_item.py @@ -1,17 +1,10 @@ -from datetime import datetime, timezone, timedelta +from datetime import timedelta from itertools import product from pathlib import PurePath import pytest from generalresearch.incite.mergers import MergeCollectionItem, MergeType -from generalresearch.incite.mergers.foundations.enriched_session import ( - EnrichedSessionMerge, -) -from generalresearch.incite.mergers.foundations.enriched_wall import ( - EnrichedWallMerge, -) -from test_utils.incite.mergers.conftest import merge_collection @pytest.mark.parametrize( diff --git a/tests/incite/mergers/test_pop_ledger.py b/tests/incite/mergers/test_pop_ledger.py index 7583faf..dc01179 100644 --- a/tests/incite/mergers/test_pop_ledger.py +++ b/tests/incite/mergers/test_pop_ledger.py @@ -1,18 +1,12 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product as iter_product -from typing import Optional import pandas as pd import pytest -from distributed.utils_test import client_no_amm from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) -from test_utils.incite.collections.conftest import ledger_collection -from test_utils.incite.conftest import incite_item_factory, mnt_filepath -from test_utils.incite.mergers.conftest import pop_ledger_merge -from test_utils.managers.ledger.conftest import create_main_accounts @pytest.mark.parametrize( diff --git a/tests/incite/mergers/test_ym_survey_merge.py b/tests/incite/mergers/test_ym_survey_merge.py index 9107f21..850df8a 100644 --- a/tests/incite/mergers/test_ym_survey_merge.py +++ b/tests/incite/mergers/test_ym_survey_merge.py @@ -1,25 +1,10 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product import pandas as pd import pytest # noinspection PyUnresolvedReferences -from distributed.utils_test import ( - cleanup, - client, - client_no_amm, - cluster_fixture, - gen_cluster, - loop, - loop_in_thread, -) - -from test_utils.incite.collections.conftest import session_collection, wall_collection -from test_utils.incite.mergers.conftest import ( - enriched_session_merge, - ym_survey_wall_merge, -) @pytest.mark.parametrize( diff --git a/tests/incite/test_collection_base.py b/tests/incite/test_collection_base.py index 5a63019..7e1577a 100644 --- a/tests/incite/test_collection_base.py +++ b/tests/incite/test_collection_base.py @@ -10,7 +10,6 @@ import pytest from _pytest._code.code import ExceptionInfo from generalresearch.incite.base import CollectionBase -from test_utils.incite.conftest import mnt_filepath AGO_15min = (datetime.now(tz=UTC) - timedelta(minutes=15)).replace(microsecond=0) AGO_1HR = (datetime.now(tz=UTC) - timedelta(hours=1)).replace(microsecond=0) @@ -108,7 +107,7 @@ class TestCollectionBase: with pytest.raises(expected_exception=ValueError) as cm: cm: ExceptionInfo - CollectionBase(offset=f"59sec", archive_path=mnt_filepath.data_src) + CollectionBase(offset="59sec", archive_path=mnt_filepath.data_src) assert "Must be equal to, or longer than 1 min" in str(cm.value) with pytest.raises(expected_exception=ValueError) as cm: diff --git a/tests/incite/test_collection_base_item.py b/tests/incite/test_collection_base_item.py index 3f4d023..7a0a581 100644 --- a/tests/incite/test_collection_base_item.py +++ b/tests/incite/test_collection_base_item.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from os.path import join as pjoin from pathlib import Path from uuid import uuid4 diff --git a/tests/incite/test_grl_flow.py b/tests/incite/test_grl_flow.py index c632f9a..6aea182 100644 --- a/tests/incite/test_grl_flow.py +++ b/tests/incite/test_grl_flow.py @@ -1,15 +1,16 @@ class TestGRLFlow: def test_init(self, mnt_filepath, thl_web_rr): + from generalresearch.incite.collections.thl_web import ( + LedgerDFCollection, + TaskAdjustmentDFCollection, + ) from generalresearch.incite.defaults import ( ledger_df_collection, task_df_collection, - pop_ledger as plm, ) - - from generalresearch.incite.collections.thl_web import ( - LedgerDFCollection, - TaskAdjustmentDFCollection, + from generalresearch.incite.defaults import ( + pop_ledger as plm, ) from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge diff --git a/tests/incite/test_interval_idx.py b/tests/incite/test_interval_idx.py index ea2bced..3034c21 100644 --- a/tests/incite/test_interval_idx.py +++ b/tests/incite/test_interval_idx.py @@ -1,5 +1,6 @@ +from datetime import datetime + import pandas as pd -from datetime import datetime, timezone, timedelta class TestIntervalIndex: diff --git a/tests/managers/gr/test_authentication.py b/tests/managers/gr/test_authentication.py index 53b6931..b9f43a6 100644 --- a/tests/managers/gr/test_authentication.py +++ b/tests/managers/gr/test_authentication.py @@ -1,11 +1,9 @@ import logging -from random import randint from uuid import uuid4 import pytest from generalresearch.models.gr.authentication import GRUser -from test_utils.models.conftest import gr_user SSO_ISSUER = "" @@ -13,7 +11,6 @@ SSO_ISSUER = "" class TestGRUserManager: def test_create(self, gr_um): - from generalresearch.models.gr.authentication import GRUser user: GRUser = gr_um.create_dummy() instance = gr_um.get_by_id(user.id) diff --git a/tests/managers/gr/test_business.py b/tests/managers/gr/test_business.py index 7eb77f8..74a5450 100644 --- a/tests/managers/gr/test_business.py +++ b/tests/managers/gr/test_business.py @@ -2,8 +2,6 @@ from uuid import uuid4 import pytest -from test_utils.models.conftest import business - class TestBusinessBankAccountManager: @@ -12,8 +10,8 @@ class TestBusinessBankAccountManager: def test_create(self, business, business_bank_account_manager): from generalresearch.models.gr.business import ( - TransferMethod, BusinessBankAccount, + TransferMethod, ) instance = business_bank_account_manager.create( diff --git a/tests/managers/gr/test_team.py b/tests/managers/gr/test_team.py index 9215da4..0918ab8 100644 --- a/tests/managers/gr/test_team.py +++ b/tests/managers/gr/test_team.py @@ -1,7 +1,5 @@ from uuid import uuid4 -from test_utils.models.conftest import team - class TestMembershipManager: diff --git a/tests/managers/leaderboard.py b/tests/managers/leaderboard.py index 149bdbb..7773030 100644 --- a/tests/managers/leaderboard.py +++ b/tests/managers/leaderboard.py @@ -1,7 +1,7 @@ import os import time import zoneinfo -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from uuid import uuid4 diff --git a/tests/managers/network/test_label.py b/tests/managers/network/test_label.py index 5b9a790..bfc7518 100644 --- a/tests/managers/network/test_label.py +++ b/tests/managers/network/test_label.py @@ -9,8 +9,8 @@ from generalresearch.managers.network.label import IPLabelManager from generalresearch.models.network.label import ( IPLabel, IPLabelKind, - IPLabelSource, IPLabelMetadata, + IPLabelSource, ) from generalresearch.models.thl.ipinfo import normalize_ip diff --git a/tests/managers/test_events.py b/tests/managers/test_events.py index 6941c00..cc9a1bf 100644 --- a/tests/managers/test_events.py +++ b/tests/managers/test_events.py @@ -1,11 +1,10 @@ import math import random import time -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from functools import partial from math import floor -from typing import Optional from uuid import uuid4 import pytest @@ -384,7 +383,7 @@ class TestTaskStatsManager: "task_created_count_last_24h": AggregateBySource(total=0), } event_manager.set_source_task_stats( - source=Source.TESTING, live_task_count=0, live_tasks_max_payout=Decimal("0") + source=Source.TESTING, live_task_count=0, live_tasks_max_payout=Decimal(0) ) assert event_manager.get_task_stats_raw() == { "live_task_count": AggregateBySource( @@ -400,7 +399,7 @@ class TestTaskStatsManager: event_manager.set_source_task_stats( source=Source.TESTING, live_task_count=0, - live_tasks_max_payout=Decimal("0"), + live_tasks_max_payout=Decimal(0), created_count=10, ) res = event_manager.get_task_stats_raw() @@ -414,7 +413,7 @@ class TestTaskStatsManager: event_manager.set_source_task_stats( source=Source.TESTING, live_task_count=0, - live_tasks_max_payout=Decimal("0"), + live_tasks_max_payout=Decimal(0), created_count=10, ) res = event_manager.get_task_stats_raw() @@ -428,7 +427,7 @@ class TestTaskStatsManager: event_manager.set_source_task_stats( source=Source.TESTING2, live_task_count=0, - live_tasks_max_payout=Decimal("0"), + live_tasks_max_payout=Decimal(0), created_count=1, ) res = event_manager.get_task_stats_raw() @@ -481,7 +480,7 @@ class TestChannelsSubscriptions: wall = Wall( req_survey_id="a", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), source=Source.TESTING, session_id=session.id, user_id=user.user_id, @@ -497,7 +496,7 @@ class TestChannelsSubscriptions: status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, finished=datetime.now(tz=UTC), - cpi=Decimal("1"), + cpi=Decimal(1), ) event_manager.handle_task_finish(wall, session, user) msg = event_subscriber.get_next_message() diff --git a/tests/managers/test_userpid.py b/tests/managers/test_userpid.py index 4a3f699..36c2de9 100644 --- a/tests/managers/test_userpid.py +++ b/tests/managers/test_userpid.py @@ -1,11 +1,10 @@ import pytest from pydantic import MySQLDsn -from generalresearch.managers.marketplace.user_pid import UserPidMultiManager -from generalresearch.sql_helper import SqlHelper from generalresearch.managers.cint.user_pid import CintUserPidManager from generalresearch.managers.dynata.user_pid import DynataUserPidManager from generalresearch.managers.innovate.user_pid import InnovateUserPidManager +from generalresearch.managers.marketplace.user_pid import UserPidMultiManager from generalresearch.managers.morning.user_pid import MorningUserPidManager # from generalresearch.managers.precision import PrecisionUserPidManager @@ -13,6 +12,7 @@ from generalresearch.managers.prodege.user_pid import ProdegeUserPidManager from generalresearch.managers.repdata.user_pid import RepdataUserPidManager from generalresearch.managers.sago.user_pid import SagoUserPidManager from generalresearch.managers.spectrum.user_pid import SpectrumUserPidManager +from generalresearch.sql_helper import SqlHelper dsn = "" diff --git a/tests/managers/thl/test_contest/test_leaderboard.py b/tests/managers/thl/test_contest/test_leaderboard.py index 1a52f83..7adea9c 100644 --- a/tests/managers/thl/test_contest/test_leaderboard.py +++ b/tests/managers/thl/test_contest/test_leaderboard.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from zoneinfo import ZoneInfo from generalresearch.currency import USDCent @@ -12,12 +12,6 @@ from generalresearch.models.thl.contest.leaderboard import ( ) from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User -from test_utils.managers.contest.conftest import ( - leaderboard_contest_create as contest_create, -) -from test_utils.managers.contest.conftest import ( - leaderboard_contest_in_db as contest_in_db, -) class TestLeaderboardContestCRUD: diff --git a/tests/managers/thl/test_contest/test_milestone.py b/tests/managers/thl/test_contest/test_milestone.py index 66c5dc4..ed0bbb5 100644 --- a/tests/managers/thl/test_contest/test_milestone.py +++ b/tests/managers/thl/test_contest/test_milestone.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from generalresearch.models.thl.contest.definitions import ( ContestEndReason, @@ -12,18 +12,6 @@ from generalresearch.models.thl.contest.milestone import ( ) from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User -from test_utils.managers.contest.conftest import ( - milestone_contest as contest, -) -from test_utils.managers.contest.conftest import ( - milestone_contest_create as contest_create, -) -from test_utils.managers.contest.conftest import ( - milestone_contest_factory as contest_factory, -) -from test_utils.managers.contest.conftest import ( - milestone_contest_in_db as contest_in_db, -) class TestMilestoneContest: diff --git a/tests/managers/thl/test_contest/test_raffle.py b/tests/managers/thl/test_contest/test_raffle.py index 5804ea3..736a5e9 100644 --- a/tests/managers/thl/test_contest/test_raffle.py +++ b/tests/managers/thl/test_contest/test_raffle.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime import pytest from pydantic import ValidationError @@ -28,18 +28,6 @@ from generalresearch.models.thl.contest.raffle import ( ) from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User -from test_utils.managers.contest.conftest import ( - raffle_contest as contest, -) -from test_utils.managers.contest.conftest import ( - raffle_contest_create as contest_create, -) -from test_utils.managers.contest.conftest import ( - raffle_contest_factory as contest_factory, -) -from test_utils.managers.contest.conftest import ( - raffle_contest_in_db as contest_in_db, -) class TestRaffleContest: diff --git a/tests/managers/thl/test_harmonized_uqa.py b/tests/managers/thl/test_harmonized_uqa.py index 3b6df48..81ac080 100644 --- a/tests/managers/thl/test_harmonized_uqa.py +++ b/tests/managers/thl/test_harmonized_uqa.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime import pytest diff --git a/tests/managers/thl/test_ipinfo.py b/tests/managers/thl/test_ipinfo.py index 847b00c..61d4d19 100644 --- a/tests/managers/thl/test_ipinfo.py +++ b/tests/managers/thl/test_ipinfo.py @@ -1,9 +1,9 @@ import faker from generalresearch.managers.thl.ipinfo import ( + GeoIpInfoManager, IPGeonameManager, IPInformationManager, - GeoIpInfoManager, ) from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation diff --git a/tests/managers/thl/test_ledger/test_lm_accounts.py b/tests/managers/thl/test_ledger/test_lm_accounts.py index faef5fb..11b2835 100644 --- a/tests/managers/thl/test_ledger/test_lm_accounts.py +++ b/tests/managers/thl/test_ledger/test_lm_accounts.py @@ -1,4 +1,3 @@ -from collections.abc import Callable from itertools import product as iproduct from random import randint from typing import TYPE_CHECKING @@ -22,7 +21,6 @@ from generalresearch.models.thl.ledger import ( if TYPE_CHECKING: from pydantic import PositiveInt - from generalresearch.config import GRLSettings from generalresearch.currency import LedgerCurrency from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager from generalresearch.models.custom_types import AccountType, Direction, UUIDStr @@ -32,10 +30,6 @@ if TYPE_CHECKING: LedgerAccount, LedgerTransaction, ) - from generalresearch.models.thl.product import Product - from generalresearch.models.thl.session import Session - from generalresearch.models.thl.user import User - from generalresearch.models.thl.wallet import PayoutType @pytest.mark.parametrize( diff --git a/tests/managers/thl/test_ledger/test_lm_tx_locks.py b/tests/managers/thl/test_ledger/test_lm_tx_locks.py index 07c3712..020b74a 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx_locks.py +++ b/tests/managers/thl/test_ledger/test_lm_tx_locks.py @@ -1,6 +1,6 @@ import logging from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal import pytest @@ -23,7 +23,6 @@ from generalresearch.models.thl.session import ( WallAdjustedStatus, ) from generalresearch.models.thl.user import User -from test_utils.models.conftest import product_user_wallet_no, session, user_factory logger = logging.getLogger("LedgerManager") @@ -56,7 +55,7 @@ class TestLedgerLocks: s1 = session_factory( user=user, wall_count=3, - wall_req_cpis=[Decimal("1.23"), Decimal("3.21"), Decimal("4")], + wall_req_cpis=[Decimal("1.23"), Decimal("3.21"), Decimal(4)], wall_statuses=[Status.COMPLETE, Status.COMPLETE, Status.COMPLETE], ) @@ -187,7 +186,7 @@ class TestLedgerLocks: user_id=user.user_id, source=Source.DYNATA, req_survey_id="xxx", - req_cpi=Decimal("4"), + req_cpi=Decimal(4), session_id=1, status=Status.COMPLETE, status_code_1=StatusCode1.COMPLETE, @@ -292,13 +291,7 @@ class TestLedgerLocks: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": session.started + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=session.started + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) print(thl_net, commission_amount, bp_pay, user_pay) @@ -336,13 +329,7 @@ class TestLedgerLocks: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": session.started + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=session.started + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) print(thl_net, commission_amount, bp_pay, user_pay) diff --git a/tests/managers/thl/test_ledger/test_thl_lm_accounts.py b/tests/managers/thl/test_ledger/test_thl_lm_accounts.py index 01d5fe1..8d7d828 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_accounts.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_accounts.py @@ -8,9 +8,9 @@ class TestThlLedgerManagerAccounts: def test_get_account_or_create_user_wallet(self, user, thl_lm, lm): from generalresearch.currency import LedgerCurrency from generalresearch.models.thl.ledger import ( - LedgerAccount, - Direction, AccountType, + Direction, + LedgerAccount, ) account = thl_lm.get_account_or_create_user_wallet(user=user) @@ -31,9 +31,9 @@ class TestThlLedgerManagerAccounts: def test_get_account_or_create_bp_wallet(self, product, thl_lm, lm): from generalresearch.currency import LedgerCurrency from generalresearch.models.thl.ledger import ( - LedgerAccount, - Direction, AccountType, + Direction, + LedgerAccount, ) account = thl_lm.get_account_or_create_bp_wallet(product=product) @@ -54,8 +54,8 @@ class TestThlLedgerManagerAccounts: def test_get_account_or_create_bp_commission(self, product, thl_lm, lm): from generalresearch.currency import LedgerCurrency from generalresearch.models.thl.ledger import ( - Direction, AccountType, + Direction, ) account = thl_lm.get_account_or_create_bp_commission(product=product) @@ -76,8 +76,8 @@ class TestThlLedgerManagerAccounts: def test_get_account_or_create_bp_expense(self, product, expense, thl_lm, lm): from generalresearch.currency import LedgerCurrency from generalresearch.models.thl.ledger import ( - Direction, AccountType, + Direction, ) account = thl_lm.get_account_or_create_bp_expense( @@ -98,8 +98,8 @@ class TestThlLedgerManagerAccounts: def test_get_or_create_bp_pending_payout_account(self, product, thl_lm, lm): from generalresearch.currency import LedgerCurrency from generalresearch.models.thl.ledger import ( - Direction, AccountType, + Direction, ) account = thl_lm.get_or_create_bp_pending_payout_account(product=product) @@ -132,8 +132,8 @@ class TestThlLedgerManagerAccounts: self, account_cash, account_revenue_task_complete, thl_lm, lm ): from generalresearch.models.thl.ledger import ( - LedgerAccount, AccountType, + LedgerAccount, ) res = thl_lm.get_account_task_complete_revenue() @@ -155,8 +155,8 @@ class TestThlLedgerManagerAccounts: def test_get_account_cash(self, account_cash, thl_lm, lm): from generalresearch.models.thl.ledger import ( - LedgerAccount, AccountType, + LedgerAccount, ) res = thl_lm.get_account_cash() @@ -167,10 +167,10 @@ class TestThlLedgerManagerAccounts: assert res.display_name == "Operating Cash Account" def test_get_accounts(self, setup_accounts, product, user_factory, thl_lm, lm, lam): - from generalresearch.models.thl.user import User from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerAccountDoesntExistError, ) + from generalresearch.models.thl.user import User user1: User = user_factory(product=product) user2: User = user_factory(product=product) @@ -276,9 +276,9 @@ class TestLedgerAccountManager: LedgerAccountDoesntExistError, ) from generalresearch.models.thl.ledger import ( - LedgerAccount, - Direction, AccountType, + Direction, + LedgerAccount, ) u = uuid4().hex @@ -311,8 +311,8 @@ class TestLedgerAccountManager: LedgerAccountDoesntExistError, ) from generalresearch.models.thl.ledger import ( - LedgerAccount, AccountType, + LedgerAccount, ) with pytest.raises(LedgerAccountDoesntExistError): @@ -326,10 +326,10 @@ class TestLedgerAccountManager: assert user.product.uuid == account.reference_uuid def test_get_many(self, product_factory, thl_lm, lm, lam, currency): - from generalresearch.models.thl.product import Product from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerAccountDoesntExistError, ) + from generalresearch.models.thl.product import Product p1: Product = product_factory() p2: Product = product_factory() @@ -388,9 +388,9 @@ class TestLedgerAccountManager: def test_create_account(self, thl_lm, lm, lam): from generalresearch.models.thl.ledger import ( - LedgerAccount, - Direction, AccountType, + Direction, + LedgerAccount, ) u = uuid4().hex diff --git a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py index 1fb9c01..cfb8f8f 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint from uuid import uuid4 @@ -78,13 +78,7 @@ class TestThlLedgerManagerBPPayout: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": now + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=now + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) thl_lm.create_tx_bp_payment(session=session, created=wall1.started) diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx.py b/tests/managers/thl/test_ledger/test_thl_lm_tx.py index be988a1..1130621 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint from uuid import uuid4 @@ -114,7 +114,7 @@ class TestThlLedgerTxManager: user_wallet_config=UserWalletConfig(amt=True, enabled=True), ) user = user_factory(product=product) - s1 = session_factory(user=user, wall_req_cpi=Decimal("1")) + s1 = session_factory(user=user, wall_req_cpi=Decimal(1)) status, status_code_1 = s1.determine_session_status() assert status == Status.COMPLETE @@ -207,9 +207,7 @@ class TestThlLedgerTxManager: # Update the finished timestamp, but nothing else. This means that # there is no financial changes needed session.update( - **{ - "finished": datetime.now(tz=UTC) + timedelta(minutes=10), - } + finished=datetime.now(tz=UTC) + timedelta(minutes=10) ) assert session.finished with caplog.at_level(logging.INFO): @@ -829,13 +827,7 @@ class TestThlLedgerTxManagerFlows: status, status_code_1 = s1.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = s1.determine_payments() s1.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": s1.started + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=s1.started + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) print(thl_net, commission_amount, bp_pay, user_pay) thl_lm.create_tx_bp_payment(session=s1, created=w1.started) @@ -892,13 +884,7 @@ class TestThlLedgerTxManagerFlows: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": session.started + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=session.started + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) print(thl_net, commission_amount, bp_pay, user_pay) @@ -930,13 +916,7 @@ class TestThlLedgerTxManagerFlows: status, status_code_1 = session.determine_session_status() # thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": session.started + timedelta(minutes=10), - "payout": Decimal("1.53"), - "user_payout": Decimal("1.53"), - } + status=status, status_code_1=status_code_1, finished=session.started + timedelta(minutes=10), payout=Decimal("1.53"), user_payout=Decimal("1.53") ) thl_lm.create_tx_bp_payment(session=session, created=wall1.started) @@ -971,13 +951,7 @@ class TestThlLedgerTxManagerFlows: status, status_code_1 = session.determine_session_status() # thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": session.started + timedelta(minutes=10), - "payout": Decimal("0.39"), - "user_payout": Decimal("0.26"), - } + status=status, status_code_1=status_code_1, finished=session.started + timedelta(minutes=10), payout=Decimal("0.39"), user_payout=Decimal("0.26") ) # with pytest.logs(logger, level=logging.WARNING) as cm: # tx = thl_lm.create_transaction_bp_payment(session, created=wall1.started) @@ -1310,13 +1284,7 @@ class TestThlLedgerManagerAdj: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": utc_hour_ago + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=utc_hour_ago + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) thl_lm.create_tx_bp_payment(session, created=wall1.started) @@ -1484,13 +1452,7 @@ class TestThlLedgerManagerAdj: assert user_pay == Decimal("1.52") session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": now + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=now + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) tx = thl_lm.create_tx_bp_adjustment(session=session, created=wall1.started) @@ -1664,13 +1626,7 @@ class TestThlLedgerManagerAdj: status, status_code_1 = session.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = session.determine_payments() session.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": utc_hour_ago + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=utc_hour_ago + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) thl_lm.create_tx_bp_payment(session, created=wall1.started) diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py index 9253ff0..cd6ea79 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from uuid import uuid4 @@ -12,7 +12,6 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( from generalresearch.models.thl.payout import UserPayoutEvent from generalresearch.models.thl.user import User from generalresearch.models.thl.wallet import PayoutType -from test_utils.managers.ledger.conftest import create_main_accounts class TestLedgerManagerAMT: diff --git a/tests/managers/thl/test_ledger/test_thl_pem.py b/tests/managers/thl/test_ledger/test_thl_pem.py index 5fb9e7d..7a4f0c9 100644 --- a/tests/managers/thl/test_ledger/test_thl_pem.py +++ b/tests/managers/thl/test_ledger/test_thl_pem.py @@ -1,6 +1,5 @@ -import uuid from random import randint -from uuid import uuid4, UUID +from uuid import UUID, uuid4 import pytest @@ -33,7 +32,6 @@ class TestThlPayoutEventManager: thl_lm, brokerage_product_payout_event_manager, ): - from generalresearch.models.thl.payout import UserPayoutEvent N_PRODUCTS = randint(3, 10) N_PAYOUT_EVENTS = randint(3, 10) @@ -73,7 +71,6 @@ class TestThlPayoutEventManager: brokerage_product_payout_event_manager, thl_lm, ): - from generalresearch.models.thl.payout import UserPayoutEvent N_PRODUCTS = randint(3, 10) N_PAYOUT_EVENTS = randint(3, 10) @@ -119,7 +116,6 @@ class TestThlPayoutEventManager: description can't be None """ from generalresearch.models.thl.payout import ( - UserPayoutEvent, PayoutType, ) @@ -174,7 +170,6 @@ class TestThlPayoutEventManager: brokerage_product_payout_event_manager, lm, ): - from generalresearch.models.thl.payout import UserPayoutEvent delete_ledger_db() create_main_accounts() diff --git a/tests/managers/thl/test_ledger/test_user_txs.py b/tests/managers/thl/test_ledger/test_user_txs.py index b4b0437..f83641e 100644 --- a/tests/managers/thl/test_ledger/test_user_txs.py +++ b/tests/managers/thl/test_ledger/test_user_txs.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime from decimal import Decimal from typing import TYPE_CHECKING from uuid import uuid4 @@ -8,7 +8,6 @@ from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerMana from generalresearch.managers.thl.user_compensate import user_compensate from generalresearch.models.thl.definitions import ( Status, - WallAdjustedStatus, ) from generalresearch.models.thl.ledger import ( TransactionType, diff --git a/tests/managers/thl/test_ledger/test_wallet.py b/tests/managers/thl/test_ledger/test_wallet.py index a0abd7c..bb49cd8 100644 --- a/tests/managers/thl/test_ledger/test_wallet.py +++ b/tests/managers/thl/test_ledger/test_wallet.py @@ -4,10 +4,10 @@ from uuid import uuid4 import pytest from generalresearch.models.thl.product import ( - UserWalletConfig, PayoutConfig, PayoutTransformation, PayoutTransformationPercentArgs, + UserWalletConfig, ) from generalresearch.models.thl.user import User diff --git a/tests/managers/thl/test_product.py b/tests/managers/thl/test_product.py index 78d5dde..f8dd44d 100644 --- a/tests/managers/thl/test_product.py +++ b/tests/managers/thl/test_product.py @@ -5,15 +5,14 @@ import pytest from generalresearch.models import Source from generalresearch.models.thl.product import ( Product, + ProfilingConfig, SourceConfig, - UserCreateConfig, SourcesConfig, - UserHealthConfig, - ProfilingConfig, - SupplyPolicy, SupplyConfig, + SupplyPolicy, + UserCreateConfig, + UserHealthConfig, ) -from test_utils.models.conftest import product_factory class TestProductManagerGetMethods: diff --git a/tests/managers/thl/test_product_prod.py b/tests/managers/thl/test_product_prod.py index 7b4f677..f5aa78d 100644 --- a/tests/managers/thl/test_product_prod.py +++ b/tests/managers/thl/test_product_prod.py @@ -4,7 +4,6 @@ from uuid import uuid4 import pytest from generalresearch.models.thl.product import Product -from test_utils.models.conftest import product_factory logger = logging.getLogger() @@ -79,4 +78,3 @@ class TestProductManagerGetAll: products = product_manager.get_all(rand_limit=50) logger.info(f"Fetching {len(products)} product uuids") # todo: once timebucks stops spamming broken accounts, fetch more - pass diff --git a/tests/managers/thl/test_profiling/test_user_upk.py b/tests/managers/thl/test_profiling/test_user_upk.py index 491e2b1..8b995b1 100644 --- a/tests/managers/thl/test_profiling/test_user_upk.py +++ b/tests/managers/thl/test_profiling/test_user_upk.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from generalresearch.managers.thl.profiling.user_upk import UserUpkManager diff --git a/tests/managers/thl/test_session_manager.py b/tests/managers/thl/test_session_manager.py index 6bedc2b..adcbe25 100644 --- a/tests/managers/thl/test_session_manager.py +++ b/tests/managers/thl/test_session_manager.py @@ -7,11 +7,10 @@ from faker import Faker from generalresearch.models import DeviceType from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.definitions import ( + SessionStatusCode2, Status, StatusCode1, - SessionStatusCode2, ) -from test_utils.models.conftest import user fake = Faker() @@ -21,8 +20,8 @@ class TestSessionManager: bucket = Bucket( loi_min=timedelta(seconds=60), loi_max=timedelta(seconds=120), - user_payout_min=Decimal("1"), - user_payout_max=Decimal("2"), + user_payout_min=Decimal(1), + user_payout_max=Decimal(2), ) s1 = session_manager.create( diff --git a/tests/managers/thl/test_survey.py b/tests/managers/thl/test_survey.py index 4b4a579..37f0b66 100644 --- a/tests/managers/thl/test_survey.py +++ b/tests/managers/thl/test_survey.py @@ -1,5 +1,5 @@ import uuid -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal import pytest diff --git a/tests/managers/thl/test_survey_penalty.py b/tests/managers/thl/test_survey_penalty.py index 4c7dc08..2a3cdc2 100644 --- a/tests/managers/thl/test_survey_penalty.py +++ b/tests/managers/thl/test_survey_penalty.py @@ -1,7 +1,6 @@ import uuid import pytest -from cachetools.keys import _HashedTuple from generalresearch.models import Source from generalresearch.models.thl.survey.penalty import ( diff --git a/tests/managers/thl/test_task_adjustment.py b/tests/managers/thl/test_task_adjustment.py index 71e3535..43337b6 100644 --- a/tests/managers/thl/test_task_adjustment.py +++ b/tests/managers/thl/test_task_adjustment.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint diff --git a/tests/managers/thl/test_task_status.py b/tests/managers/thl/test_task_status.py index 468fd5e..93a624d 100644 --- a/tests/managers/thl/test_task_status.py +++ b/tests/managers/thl/test_task_status.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal import pytest diff --git a/tests/managers/thl/test_user_manager/test_base.py b/tests/managers/thl/test_user_manager/test_base.py index 2704490..7d83c11 100644 --- a/tests/managers/thl/test_user_manager/test_base.py +++ b/tests/managers/thl/test_user_manager/test_base.py @@ -1,5 +1,5 @@ import logging -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from random import randint from uuid import uuid4 @@ -156,7 +156,7 @@ class TestCreateUserManager: # make sure thl_user row is created res_thl_user = thl_web_rw.execute_sql_query( - query=f""" + query=""" SELECT * FROM thl_user AS u WHERE u.id = %s diff --git a/tests/managers/thl/test_user_manager/test_mysql.py b/tests/managers/thl/test_user_manager/test_mysql.py index 0313bbf..d414a13 100644 --- a/tests/managers/thl/test_user_manager/test_mysql.py +++ b/tests/managers/thl/test_user_manager/test_mysql.py @@ -1,4 +1,3 @@ -from test_utils.models.conftest import user, user_manager class TestUserManagerMysqlNew: diff --git a/tests/managers/thl/test_user_manager/test_user_fetch.py b/tests/managers/thl/test_user_manager/test_user_fetch.py index a4b3d57..7c9e012 100644 --- a/tests/managers/thl/test_user_manager/test_user_fetch.py +++ b/tests/managers/thl/test_user_manager/test_user_fetch.py @@ -3,7 +3,6 @@ from uuid import uuid4 import pytest from generalresearch.models.thl.user import User -from test_utils.models.conftest import product, user_manager, user_factory class TestUserManagerFetch: diff --git a/tests/managers/thl/test_user_manager/test_user_metadata.py b/tests/managers/thl/test_user_manager/test_user_metadata.py index 91dc16a..19b3d9f 100644 --- a/tests/managers/thl/test_user_manager/test_user_metadata.py +++ b/tests/managers/thl/test_user_manager/test_user_metadata.py @@ -3,7 +3,6 @@ from uuid import uuid4 import pytest from generalresearch.models.thl.user_profile import UserMetadata -from test_utils.models.conftest import user, user_manager, user_factory class TestUserMetadataManager: diff --git a/tests/managers/thl/test_user_streak.py b/tests/managers/thl/test_user_streak.py index ef25e2b..be0729c 100644 --- a/tests/managers/thl/test_user_streak.py +++ b/tests/managers/thl/test_user_streak.py @@ -1,5 +1,5 @@ import copy -from datetime import UTC, date, datetime, timedelta, timezone +from datetime import UTC, date, datetime, timedelta from decimal import Decimal from zoneinfo import ZoneInfo diff --git a/tests/managers/thl/test_userhealth.py b/tests/managers/thl/test_userhealth.py index e2ea4a8..be8f9c8 100644 --- a/tests/managers/thl/test_userhealth.py +++ b/tests/managers/thl/test_userhealth.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 import faker diff --git a/tests/managers/thl/test_wall_manager.py b/tests/managers/thl/test_wall_manager.py index 46d199f..5abc648 100644 --- a/tests/managers/thl/test_wall_manager.py +++ b/tests/managers/thl/test_wall_manager.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from uuid import uuid4 @@ -10,7 +10,6 @@ from generalresearch.models.thl.session import ( Status, StatusCode1, ) -from test_utils.models.conftest import session, user class TestWallManager: @@ -92,7 +91,7 @@ class TestWallManager: source=Source.DYNATA, buyer_id="123", req_survey_id="456", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), ) assert w is not None @@ -110,7 +109,7 @@ class TestWallManager: source=Source.DYNATA, buyer_id="123", req_survey_id="456", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), ) wall_manager.report( wall=w1, @@ -151,7 +150,7 @@ class TestWallManager: source=Source.DYNATA, buyer_id="123", req_survey_id="456", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), ) finish_ts = utc_hour_ago + timedelta(minutes=10) @@ -190,7 +189,7 @@ class TestWallManager: source=Source.DYNATA, buyer_id="123", req_survey_id="456", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), ) res = wall_manager.filter_wall_attempts(user_id=user.user_id) assert len(res) == 1 @@ -202,7 +201,7 @@ class TestWallManager: source=Source.DYNATA, buyer_id="123", req_survey_id="555", - req_cpi=Decimal("1"), + req_cpi=Decimal(1), ) res = wall_manager.filter_wall_attempts(user_id=user.user_id) assert len(res) == 2 diff --git a/tests/models/admin/test_report_request.py b/tests/models/admin/test_report_request.py index 4626ab4..5b2ff0d 100644 --- a/tests/models/admin/test_report_request.py +++ b/tests/models/admin/test_report_request.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime import pandas as pd import pytest diff --git a/tests/models/custom_types/test_aware_datetime.py b/tests/models/custom_types/test_aware_datetime.py index 043fba0..7c45710 100644 --- a/tests/models/custom_types/test_aware_datetime.py +++ b/tests/models/custom_types/test_aware_datetime.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime import pytest import pytz diff --git a/tests/models/custom_types/test_dsn.py b/tests/models/custom_types/test_dsn.py index 050976e..eb27526 100644 --- a/tests/models/custom_types/test_dsn.py +++ b/tests/models/custom_types/test_dsn.py @@ -1,4 +1,3 @@ -from typing import Optional from uuid import uuid4 import pytest diff --git a/tests/models/dynata/test_eligbility.py b/tests/models/dynata/test_eligbility.py index 16cad26..27de5b3 100644 --- a/tests/models/dynata/test_eligbility.py +++ b/tests/models/dynata/test_eligbility.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime class TestEligibility: diff --git a/tests/models/gr/test_authentication.py b/tests/models/gr/test_authentication.py index 51595a7..ac39e64 100644 --- a/tests/models/gr/test_authentication.py +++ b/tests/models/gr/test_authentication.py @@ -2,7 +2,7 @@ import binascii import json import os from collections.abc import Callable -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from random import randint from uuid import uuid4 diff --git a/tests/models/gr/test_business.py b/tests/models/gr/test_business.py index 716ec75..948acb3 100644 --- a/tests/models/gr/test_business.py +++ b/tests/models/gr/test_business.py @@ -1,7 +1,6 @@ import os -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal -from typing import Optional from uuid import uuid4 import pandas as pd diff --git a/tests/models/innovate/test_question.py b/tests/models/innovate/test_question.py index 330f919..b0c2964 100644 --- a/tests/models/innovate/test_question.py +++ b/tests/models/innovate/test_question.py @@ -1,15 +1,15 @@ from generalresearch.models import Source from generalresearch.models.innovate.question import ( InnovateQuestion, - InnovateQuestionType, InnovateQuestionOption, + InnovateQuestionType, ) from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestionSelectorTE, UpkQuestion, + UpkQuestionChoice, UpkQuestionSelectorMC, + UpkQuestionSelectorTE, UpkQuestionType, - UpkQuestionChoice, ) diff --git a/tests/models/legacy/test_user_question_answer_in.py b/tests/models/legacy/test_user_question_answer_in.py index 224334a..1904798 100644 --- a/tests/models/legacy/test_user_question_answer_in.py +++ b/tests/models/legacy/test_user_question_answer_in.py @@ -263,12 +263,12 @@ class TestUserQuestionAnswerIn: UserQuestionAnswerIn, ) - for qid in { + for qid in ( "2fbedb2b9f7647b09ff5e52fa119cc5e", "4030c52371b04e80b64e058d9c5b82e9", "a91cb1dea814480dba12d9b7b48696dd", "1d1e2e8380ac474b87fb4e4c569b48df", - }: + ): # This is the UserAgent question which only allows a single answer with pytest.raises(ValueError) as cm: UserQuestionAnswerIn.model_validate( diff --git a/tests/models/morning/test.py b/tests/models/morning/test.py index 222cb93..7474766 100644 --- a/tests/models/morning/test.py +++ b/tests/models/morning/test.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from generalresearch.models.morning.question import MorningQuestion diff --git a/tests/models/network/test_mtr.py b/tests/models/network/test_mtr.py index 2965300..840a773 100644 --- a/tests/models/network/test_mtr.py +++ b/tests/models/network/test_mtr.py @@ -1,7 +1,7 @@ -from generalresearch.models.network.mtr.execute import execute_mtr import faker -from generalresearch.models.network.tool_run import ToolName, ToolClass +from generalresearch.models.network.mtr.execute import execute_mtr +from generalresearch.models.network.tool_run import ToolClass, ToolName fake = faker.Faker() diff --git a/tests/models/network/test_nmap_parser.py b/tests/models/network/test_nmap_parser.py index abc83c9..7822380 100644 --- a/tests/models/network/test_nmap_parser.py +++ b/tests/models/network/test_nmap_parser.py @@ -4,6 +4,7 @@ import pytest from generalresearch.models.network.nmap.parser import parse_nmap_xml + @pytest.fixture def nmap_raw_output_2(request) -> str: fp = os.path.join(request.config.rootpath, "data/nmaprun2.xml") diff --git a/tests/models/prodege/test_survey_participation.py b/tests/models/prodege/test_survey_participation.py index 3b35d0c..e1ba9ab 100644 --- a/tests/models/prodege/test_survey_participation.py +++ b/tests/models/prodege/test_survey_participation.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta class TestProdegeParticipation: diff --git a/tests/models/spectrum/test_question.py b/tests/models/spectrum/test_question.py index 4f92961..57d260d 100644 --- a/tests/models/spectrum/test_question.py +++ b/tests/models/spectrum/test_question.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from generalresearch.models import Source from generalresearch.models.spectrum.question import ( diff --git a/tests/models/spectrum/test_survey.py b/tests/models/spectrum/test_survey.py index 5e095a3..7365c7e 100644 --- a/tests/models/spectrum/test_survey.py +++ b/tests/models/spectrum/test_survey.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal diff --git a/tests/models/spectrum/test_survey_manager.py b/tests/models/spectrum/test_survey_manager.py index 11970bf..ce26c44 100644 --- a/tests/models/spectrum/test_survey_manager.py +++ b/tests/models/spectrum/test_survey_manager.py @@ -1,6 +1,6 @@ import copy import logging -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from decimal import Decimal from pymysql import IntegrityError diff --git a/tests/models/test_finance.py b/tests/models/test_finance.py index 6dcd441..3d3ff3a 100644 --- a/tests/models/test_finance.py +++ b/tests/models/test_finance.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from itertools import product as iter_product from random import randint from uuid import uuid4 diff --git a/tests/models/thl/question/test_question_info.py b/tests/models/thl/question/test_question_info.py index 945ee7a..b619fc3 100644 --- a/tests/models/thl/question/test_question_info.py +++ b/tests/models/thl/question/test_question_info.py @@ -1,6 +1,6 @@ from generalresearch.models.thl.profiling.upk_property import ( - UpkProperty, ProfilingInfo, + UpkProperty, ) diff --git a/tests/models/thl/test_adjustments.py b/tests/models/thl/test_adjustments.py index c2c035d..30e9bce 100644 --- a/tests/models/thl/test_adjustments.py +++ b/tests/models/thl/test_adjustments.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal import pytest @@ -426,7 +426,7 @@ class TestAdjustments: s1 = session_factory( user=user, wall_count=1, - wall_req_cpi=Decimal("1"), + wall_req_cpi=Decimal(1), final_status=Status.COMPLETE, started=utc_hour_ago, ) @@ -525,7 +525,7 @@ class TestAdjustments: s1 = session_factory( user=user, wall_count=1, - wall_req_cpi=Decimal("1"), + wall_req_cpi=Decimal(1), final_status=Status.COMPLETE, started=utc_hour_ago, ) @@ -534,13 +534,7 @@ class TestAdjustments: status, status_code_1 = s1.determine_session_status() thl_net, commission_amount, bp_pay, user_pay = s1.determine_payments() s1.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": utc_hour_ago + timedelta(minutes=10), - "payout": bp_pay, - "user_payout": user_pay, - } + status=status, status_code_1=status_code_1, finished=utc_hour_ago + timedelta(minutes=10), payout=bp_pay, user_payout=user_pay ) w1.update( adjusted_status=WallAdjustedStatus.ADJUSTED_TO_FAIL, @@ -592,13 +586,7 @@ class TestAdjustments: thl_net = Decimal(sum(w.cpi for w in s1.wall_events if w.is_visible_complete())) payout = user.product.determine_bp_payment(thl_net=thl_net) s1.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": utc_hour_ago + timedelta(minutes=25), - "payout": payout, - "user_payout": None, - } + status=status, status_code_1=status_code_1, finished=utc_hour_ago + timedelta(minutes=25), payout=payout, user_payout=None ) # Test. Adjust first fail to complete. Now we have 2 completes. @@ -646,13 +634,7 @@ class TestAdjustments: thl_net = Decimal(sum(w.cpi for w in s1.wall_events if w.is_visible_complete())) payout = user.product.determine_bp_payment(thl_net) s1.update( - **{ - "status": status, - "status_code_1": status_code_1, - "finished": utc_hour_ago + timedelta(minutes=25), - "payout": payout, - "user_payout": None, - } + status=status, status_code_1=status_code_1, finished=utc_hour_ago + timedelta(minutes=25), payout=payout, user_payout=None ) # Test. Adjust complete to fail. Now we have 2 fails. diff --git a/tests/models/thl/test_contest/test_leaderboard_contest.py b/tests/models/thl/test_contest/test_leaderboard_contest.py index 3efcf2f..52a4bec 100644 --- a/tests/models/thl/test_contest/test_leaderboard_contest.py +++ b/tests/models/thl/test_contest/test_leaderboard_contest.py @@ -1,4 +1,4 @@ -from datetime import UTC, timezone +from datetime import UTC from uuid import uuid4 import pytest diff --git a/tests/models/thl/test_ledger.py b/tests/models/thl/test_ledger.py index 5edcc9d..7066180 100644 --- a/tests/models/thl/test_ledger.py +++ b/tests/models/thl/test_ledger.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timezone +from datetime import UTC, datetime from uuid import uuid4 import pytest diff --git a/tests/models/thl/test_payout.py b/tests/models/thl/test_payout.py index 7068a41..d0b0acc 100644 --- a/tests/models/thl/test_payout.py +++ b/tests/models/thl/test_payout.py @@ -5,14 +5,13 @@ from pydantic import ValidationError from generalresearch.currency import USDCent from generalresearch.models.gr import Team +from generalresearch.models.gr.business import Business, BusinessAddress, BusinessType from generalresearch.models.thl.payout import ( - BusinessPayoutEvent, BrokerageProductPayoutEvent, + BusinessPayoutEvent, ) from generalresearch.models.thl.wallet import PayoutType -from generalresearch.models.gr.business import Business, BusinessAddress, BusinessType - class TestBusinessPayoutEvent: diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index 78bc10a..bc95c2d 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -3,7 +3,7 @@ from __future__ import annotations import os import shutil from collections.abc import Callable -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from uuid import uuid4 diff --git a/tests/models/thl/test_upkquestion.py b/tests/models/thl/test_upkquestion.py index d32875c..99d7871 100644 --- a/tests/models/thl/test_upkquestion.py +++ b/tests/models/thl/test_upkquestion.py @@ -201,18 +201,10 @@ class TestUpkQuestion: ) q = MorningQuestion( - **{ - "id": "gender", - "country_iso": "us", - "language_iso": "eng", - "name": "Gender", - "text": "What is your gender?", - "type": "s", - "options": [ + id="gender", country_iso="us", language_iso="eng", name="Gender", text="What is your gender?", type="s", options=[ {"id": "1", "text": "yes", "order": 1}, {"id": "2", "text": "no", "order": 2}, - ], - } + ] ) q.to_upk_question() q = MorningQuestion( diff --git a/tests/models/thl/test_user.py b/tests/models/thl/test_user.py index a4f331a..aafce68 100644 --- a/tests/models/thl/test_user.py +++ b/tests/models/thl/test_user.py @@ -197,7 +197,7 @@ class TestUserProductUserID: assert "Input should be a valid string" in str(cm.value) with pytest.raises(ValueError) as cm: - User(user_id=self.user_id, product_user_id=Decimal("0")) + User(user_id=self.user_id, product_user_id=Decimal(0)) assert "1 validation error for User" in str(cm.value) assert "Input should be a valid string" in str(cm.value) @@ -236,7 +236,7 @@ class TestUserProductUserID: def test_invalid_chars_slash(self): from generalresearch.models.thl.user import User - product_user_id = f"{self.randomword(50)}\{self.randomword(50)}" + product_user_id = rf"{self.randomword(50)}\{self.randomword(50)}" with pytest.raises(expected_exception=ValueError) as cm: User(user_id=self.user_id, product_user_id=product_user_id) assert "1 validation error for User" in str(cm.value) @@ -310,7 +310,7 @@ class TestUserUUID: assert "Input should be a valid string" in str(cm.value) with pytest.raises(ValueError) as cm: - User(user_id=self.user_id, uuid=Decimal("0")) + User(user_id=self.user_id, uuid=Decimal(0)) assert "1 validation error for User" in str(cm.value) assert "Input should be a valid string" in str(cm.value) diff --git a/tests/models/thl/test_user_iphistory.py b/tests/models/thl/test_user_iphistory.py index 0f050b0..d6ade9d 100644 --- a/tests/models/thl/test_user_iphistory.py +++ b/tests/models/thl/test_user_iphistory.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from generalresearch.models.thl.user_iphistory import ( UserIPHistory, diff --git a/tests/models/thl/test_user_streak.py b/tests/models/thl/test_user_streak.py index 0cacd3e..26c5e25 100644 --- a/tests/models/thl/test_user_streak.py +++ b/tests/models/thl/test_user_streak.py @@ -1,8 +1,8 @@ from datetime import datetime, timedelta +from zoneinfo import ZoneInfo import pytest from pydantic import ValidationError -from zoneinfo import ZoneInfo from generalresearch.models.thl.user_streak import ( StreakFulfillment, diff --git a/tests/models/thl/test_wall.py b/tests/models/thl/test_wall.py index 9e9483b..88914ac 100644 --- a/tests/models/thl/test_wall.py +++ b/tests/models/thl/test_wall.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal from uuid import uuid4 diff --git a/tests/models/thl/test_wall_session.py b/tests/models/thl/test_wall_session.py index 10f3cba..b39ad31 100644 --- a/tests/models/thl/test_wall_session.py +++ b/tests/models/thl/test_wall_session.py @@ -1,4 +1,4 @@ -from datetime import UTC, datetime, timedelta, timezone +from datetime import UTC, datetime, timedelta from decimal import Decimal import pytest diff --git a/tests/sql_helper.py b/tests/sql_helper.py index c4cc2ca..8ab7bdb 100644 --- a/tests/sql_helper.py +++ b/tests/sql_helper.py @@ -19,7 +19,7 @@ class TestSqlHelper: def test_scheme(self): from generalresearch.sql_helper import SqlHelper - dsn = MySQLDsn(f"mysql://root@localhost/test") + dsn = MySQLDsn("mysql://root@localhost/test") instance = SqlHelper(dsn=dsn) assert instance.is_mysql() @@ -30,7 +30,7 @@ class TestSqlHelper: # self.assertTrue(instance.is_postgresql()) with pytest.raises(ValidationError): - SqlHelper(dsn=MariaDBDsn(f"maria://root@localhost/test")) + SqlHelper(dsn=MariaDBDsn("maria://root@localhost/test")) def test_row_decode(self): from generalresearch.sql_helper import decode_uuids -- cgit v1.2.3 From 89ed44f466dc9a93d6f85931fb6eea0e9cbd27f6 Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Tue, 1 Sep 2026 10:43:29 -0700 Subject: working through circular imports, removing __init__ loaded Types / Definitions --- generalresearch/grliq/managers/event_plotter.py | 5 +- generalresearch/grliq/managers/forensic_events.py | 8 +- generalresearch/grliq/managers/forensic_results.py | 8 +- generalresearch/grliq/managers/forensic_summary.py | 16 +-- .../mergers/foundations/enriched_task_adjust.py | 2 +- .../schemas/mergers/foundations/enriched_wall.py | 2 +- .../incite/schemas/mergers/ym_wall_summary.py | 2 +- generalresearch/incite/schemas/thl_web.py | 2 +- generalresearch/managers/cint/user_pid.py | 2 +- generalresearch/managers/dynata/user_pid.py | 2 +- generalresearch/managers/events.py | 2 +- generalresearch/managers/innovate/user_pid.py | 2 +- generalresearch/managers/marketplace/user_pid.py | 2 +- generalresearch/managers/morning/user_pid.py | 2 +- generalresearch/managers/network/label.py | 6 +- generalresearch/managers/precision/user_pid.py | 2 +- generalresearch/managers/prodege/user_pid.py | 2 +- generalresearch/managers/repdata/user_pid.py | 2 +- generalresearch/managers/sago/user_pid.py | 2 +- generalresearch/managers/spectrum/user_pid.py | 2 +- generalresearch/managers/thl/buyer.py | 2 +- generalresearch/managers/thl/cashout_method.py | 2 +- generalresearch/managers/thl/contest_manager.py | 4 +- .../managers/thl/ledger_manager/ledger.py | 2 +- .../managers/thl/ledger_manager/thl_ledger.py | 2 +- generalresearch/managers/thl/payout.py | 2 +- generalresearch/managers/thl/session.py | 4 +- generalresearch/managers/thl/survey.py | 2 +- generalresearch/managers/thl/task_adjustment.py | 2 +- generalresearch/managers/thl/wall.py | 4 +- generalresearch/managers/thl/wallet/__init__.py | 2 +- generalresearch/managers/utils.py | 16 +++ generalresearch/models/cint/question.py | 2 +- generalresearch/models/cint/survey.py | 2 +- generalresearch/models/custom_types.py | 2 +- generalresearch/models/definitions.py | 114 +++++++++++++++++++++ generalresearch/models/device.py | 2 +- generalresearch/models/dynata/question.py | 2 +- generalresearch/models/dynata/survey.py | 4 +- generalresearch/models/dynata/task_collection.py | 2 +- generalresearch/models/events.py | 2 +- generalresearch/models/gr/business.py | 22 ++-- generalresearch/models/innovate/question.py | 2 +- generalresearch/models/innovate/survey.py | 8 +- generalresearch/models/legacy/bucket.py | 2 +- generalresearch/models/legacy/questions.py | 2 +- generalresearch/models/lucid/question.py | 2 +- generalresearch/models/lucid/survey.py | 2 +- generalresearch/models/morning/question.py | 2 +- generalresearch/models/morning/survey.py | 2 +- generalresearch/models/pollfish/question.py | 2 +- generalresearch/models/precision/question.py | 2 +- generalresearch/models/precision/survey.py | 2 +- generalresearch/models/prodege/question.py | 2 +- generalresearch/models/prodege/survey.py | 6 +- generalresearch/models/repdata/question.py | 2 +- generalresearch/models/repdata/survey.py | 2 +- generalresearch/models/repdata/task_collection.py | 2 +- generalresearch/models/sago/question.py | 2 +- generalresearch/models/sago/survey.py | 2 +- generalresearch/models/spectrum/question.py | 2 +- generalresearch/models/spectrum/survey.py | 2 +- generalresearch/models/spectrum/task_collection.py | 2 +- generalresearch/models/thl/__init__.py | 4 +- generalresearch/models/thl/category.py | 5 +- .../models/thl/contest/contest_entry.py | 13 +-- generalresearch/models/thl/contest/raffle.py | 4 +- generalresearch/models/thl/finance.py | 15 +-- generalresearch/models/thl/ledger.py | 2 +- generalresearch/models/thl/offerwall/__init__.py | 2 +- generalresearch/models/thl/offerwall/base.py | 2 +- generalresearch/models/thl/offerwall/cache.py | 2 +- generalresearch/models/thl/payout.py | 2 +- generalresearch/models/thl/product.py | 4 +- .../models/thl/profiling/marketplace.py | 4 +- .../models/thl/profiling/upk_question.py | 8 +- .../models/thl/profiling/upk_question_answer.py | 2 +- generalresearch/models/thl/profiling/user_info.py | 2 +- .../models/thl/profiling/user_question_answer.py | 8 +- generalresearch/models/thl/session.py | 4 +- generalresearch/models/thl/soft_pair.py | 2 +- generalresearch/models/thl/survey/__init__.py | 2 +- generalresearch/models/thl/survey/buyer.py | 2 +- generalresearch/models/thl/survey/condition.py | 2 +- generalresearch/models/thl/survey/model.py | 2 +- generalresearch/models/thl/survey/penalty.py | 2 +- generalresearch/models/thl/task_adjustment.py | 4 +- generalresearch/models/thl/user.py | 4 +- generalresearch/models/thl/user_profile.py | 2 +- generalresearch/models/thl/user_quality_event.py | 2 +- generalresearch/models/thl/user_streak.py | 2 +- .../models/thl/wallet/cashout_method.py | 4 +- generalresearch/models/thl/wallet/definitions.py | 87 ++++++++++++++++ generalresearch/models/thl/wallet/payout.py | 2 +- generalresearch/schemas/survey_stats.py | 2 +- generalresearch/wall_status_codes/__init__.py | 2 +- test_utils/conftest.py | 17 +-- test_utils/grliq/conftest.py | 12 ++- test_utils/incite/collections/conftest.py | 4 +- test_utils/incite/mergers/conftest.py | 55 +++++----- test_utils/managers/cashout_methods.py | 2 +- test_utils/managers/conftest.py | 55 +++++----- test_utils/managers/contest/conftest.py | 6 +- test_utils/managers/gr/conftest.py | 7 +- test_utils/managers/ledger/conftest.py | 14 ++- test_utils/managers/thl/conftest.py | 51 ++++----- test_utils/managers/upk/conftest.py | 9 +- test_utils/models/conftest.py | 4 +- test_utils/models/contest/conftest.py | 29 +++--- test_utils/models/gr/conftest.py | 38 +++---- test_utils/models/ledger/conftest.py | 4 +- test_utils/models/network/conftest.py | 5 +- test_utils/models/thl/conftest.py | 80 ++++++++------- test_utils/models/upk/conftest.py | 3 +- test_utils/spectrum/conftest.py | 8 +- .../incite/collections/test_df_collection_base.py | 2 +- .../collections/test_df_collection_item_base.py | 5 +- .../collections/test_df_collection_item_thl_web.py | 20 ++-- .../test_df_collection_thl_marketplaces.py | 9 +- .../collections/test_df_collection_thl_web.py | 2 +- .../mergers/foundations/test_enriched_session.py | 31 +++--- .../foundations/test_enriched_task_adjust.py | 30 +++--- .../mergers/foundations/test_enriched_wall.py | 29 +++--- .../mergers/foundations/test_user_id_product.py | 9 +- tests/incite/mergers/test_merge_collection.py | 7 +- tests/incite/mergers/test_merge_collection_item.py | 13 ++- tests/incite/mergers/test_pop_ledger.py | 21 ++-- tests/incite/mergers/test_ym_survey_merge.py | 24 +++-- tests/incite/test_collection_base.py | 6 +- tests/incite/test_collection_base_item.py | 6 +- tests/managers/gr/test_business.py | 19 ++-- tests/managers/gr/test_team.py | 15 +-- tests/managers/leaderboard.py | 5 +- tests/managers/network/test_label.py | 17 ++- tests/managers/test_events.py | 12 ++- tests/managers/test_lucid.py | 6 +- tests/managers/thl/test_buyer.py | 7 +- tests/managers/thl/test_cashout_method.py | 19 ++-- tests/managers/thl/test_category.py | 7 +- .../managers/thl/test_contest/test_leaderboard.py | 19 ++-- tests/managers/thl/test_contest/test_milestone.py | 20 ++-- tests/managers/thl/test_contest/test_raffle.py | 21 ++-- tests/managers/thl/test_harmonized_uqa.py | 7 +- tests/managers/thl/test_ipinfo.py | 7 +- tests/managers/thl/test_ledger/test_lm_accounts.py | 10 +- tests/managers/thl/test_ledger/test_lm_tx.py | 7 +- .../managers/thl/test_ledger/test_lm_tx_entries.py | 10 +- tests/managers/thl/test_ledger/test_lm_tx_locks.py | 15 +-- .../thl/test_ledger/test_lm_tx_metadata.py | 11 +- .../thl/test_ledger/test_thl_lm_accounts.py | 15 +-- .../thl/test_ledger/test_thl_lm_bp_payout.py | 22 ++-- tests/managers/thl/test_ledger/test_thl_lm_tx.py | 28 +++-- .../test_ledger/test_thl_lm_tx__user_payouts.py | 13 ++- tests/managers/thl/test_ledger/test_thl_pem.py | 21 ++-- tests/managers/thl/test_ledger/test_user_txs.py | 19 ++-- tests/managers/thl/test_ledger/test_wallet.py | 9 +- tests/managers/thl/test_payout.py | 53 +++++----- tests/managers/thl/test_product.py | 9 +- tests/managers/thl/test_product_prod.py | 5 +- tests/managers/thl/test_profiling/test_question.py | 7 +- tests/managers/thl/test_profiling/test_schema.py | 9 +- tests/managers/thl/test_profiling/test_user_upk.py | 6 +- tests/managers/thl/test_session_manager.py | 15 +-- tests/managers/thl/test_survey.py | 18 ++-- tests/managers/thl/test_survey_penalty.py | 7 +- tests/managers/thl/test_task_adjustment.py | 21 ++-- tests/managers/thl/test_task_status.py | 18 ++-- tests/managers/thl/test_user_manager/test_base.py | 18 ++-- tests/managers/thl/test_user_manager/test_mysql.py | 11 +- tests/managers/thl/test_user_manager/test_redis.py | 10 +- .../thl/test_user_manager/test_user_fetch.py | 8 +- .../thl/test_user_manager/test_user_metadata.py | 13 ++- tests/managers/thl/test_user_streak.py | 11 +- tests/managers/thl/test_userhealth.py | 19 +++- tests/managers/thl/test_wall_manager.py | 13 ++- tests/models/custom_types/test_aware_datetime.py | 4 +- tests/models/custom_types/test_dsn.py | 4 +- tests/models/custom_types/test_uuid_str.py | 4 +- tests/models/dynata/test_survey.py | 3 + tests/models/gr/test_authentication.py | 14 ++- tests/models/gr/test_base.py | 4 +- tests/models/gr/test_business.py | 40 ++++---- tests/models/gr/test_team.py | 38 ++++--- tests/models/innovate/test_question.py | 2 +- .../models/legacy/test_offerwall_parse_response.py | 2 +- .../models/legacy/test_user_question_answer_in.py | 9 +- tests/models/network/test_mtr.py | 6 +- tests/models/network/test_nmap.py | 8 +- tests/models/network/test_nmap_parser.py | 9 +- tests/models/network/test_rdns.py | 6 +- tests/models/spectrum/test_question.py | 2 +- tests/models/spectrum/test_survey.py | 2 +- tests/models/spectrum/test_survey_manager.py | 12 ++- tests/models/test_device.py | 2 +- tests/models/test_finance.py | 19 ++-- tests/models/thl/test_adjustments.py | 17 +-- tests/models/thl/test_buyer.py | 2 +- tests/models/thl/test_contest/test_contest.py | 6 +- .../thl/test_contest/test_leaderboard_contest.py | 7 +- .../models/thl/test_contest/test_raffle_contest.py | 7 +- tests/models/thl/test_marketplace_condition.py | 6 +- tests/models/thl/test_payout.py | 8 +- tests/models/thl/test_payout_format.py | 8 +- tests/models/thl/test_product.py | 37 ++++--- tests/models/thl/test_product_userwalletconfig.py | 2 +- tests/models/thl/test_soft_pair.py | 2 +- tests/models/thl/test_user.py | 9 +- tests/models/thl/test_user_metadata.py | 2 +- tests/models/thl/test_wall.py | 2 +- tests/models/thl/test_wall_session.py | 2 +- tests/test_postgres.py | 5 +- 211 files changed, 1318 insertions(+), 772 deletions(-) create mode 100644 generalresearch/managers/utils.py create mode 100644 generalresearch/models/definitions.py create mode 100644 generalresearch/models/thl/wallet/definitions.py (limited to 'test_utils/models/network') diff --git a/generalresearch/grliq/managers/event_plotter.py b/generalresearch/grliq/managers/event_plotter.py index 94b70ef..61cc52c 100644 --- a/generalresearch/grliq/managers/event_plotter.py +++ b/generalresearch/grliq/managers/event_plotter.py @@ -1,12 +1,15 @@ import html import webbrowser +from typing import TYPE_CHECKING import numpy as np from more_itertools import windowed from scipy.spatial.distance import euclidean from generalresearch.grliq.managers.colormap import turbo_colormap_data -from generalresearch.grliq.models.events import KeyboardEvent, MouseEvent + +if TYPE_CHECKING: + from generalresearch.grliq.models.events import KeyboardEvent, MouseEvent def make_events_svg( diff --git a/generalresearch/grliq/managers/forensic_events.py b/generalresearch/grliq/managers/forensic_events.py index 93da481..a97a9c2 100644 --- a/generalresearch/grliq/managers/forensic_events.py +++ b/generalresearch/grliq/managers/forensic_events.py @@ -1,7 +1,7 @@ import json from collections.abc import Collection from datetime import datetime -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 from psycopg import sql @@ -14,8 +14,10 @@ from generalresearch.grliq.models.events import ( PointerMove, TimingData, ) -from generalresearch.models.custom_types import UUIDStr -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.models.custom_types import UUIDStr + from generalresearch.pg_helper import PostgresConfig class GrlIqEventManager: diff --git a/generalresearch/grliq/managers/forensic_results.py b/generalresearch/grliq/managers/forensic_results.py index 158e582..93b1cdc 100644 --- a/generalresearch/grliq/managers/forensic_results.py +++ b/generalresearch/grliq/managers/forensic_results.py @@ -1,14 +1,16 @@ from collections.abc import Collection from datetime import datetime -from typing import Any +from typing import TYPE_CHECKING, Any from generalresearch.grliq.models.forensic_result import ( GrlIqForensicCategoryResult, Phase, ) from generalresearch.grliq.models.useragents import GrlUserAgent -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig class GrlIqCategoryResultsReader: diff --git a/generalresearch/grliq/managers/forensic_summary.py b/generalresearch/grliq/managers/forensic_summary.py index b86e1f5..c222075 100644 --- a/generalresearch/grliq/managers/forensic_summary.py +++ b/generalresearch/grliq/managers/forensic_summary.py @@ -3,14 +3,10 @@ from __future__ import annotations import statistics from collections import defaultdict from datetime import UTC, datetime, timedelta -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np -from generalresearch.grliq.managers.forensic_data import GrlIqDataManager -from generalresearch.grliq.managers.forensic_events import ( - GrlIqEventManager, -) from generalresearch.grliq.models.forensic_result import ( GrlIqCheckerResults, GrlIqForensicCategoryResult, @@ -22,8 +18,14 @@ from generalresearch.grliq.models.forensic_summary import ( TimingDataCountrySummary, UserForensicSummary, ) -from generalresearch.models.thl.user import User -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.grliq.managers.forensic_data import GrlIqDataManager + from generalresearch.grliq.managers.forensic_events import ( + GrlIqEventManager, + ) + from generalresearch.models.thl.user import User + from generalresearch.redis_helper import RedisConfig def calculate_category_summary( diff --git a/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py b/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py index ead42d9..ac9a35a 100644 --- a/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py +++ b/generalresearch/incite/schemas/mergers/foundations/enriched_task_adjust.py @@ -4,7 +4,7 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.incite.schemas import ARCHIVE_AFTER, ORDER_KEY from generalresearch.incite.schemas.thl_web import THLTaskAdjustmentSchema from generalresearch.locales import Localelator -from generalresearch.models import DeviceType, Source +from generalresearch.models.definitions import DeviceType, Source from generalresearch.models.thl.definitions import ( WallAdjustedStatus, ) diff --git a/generalresearch/incite/schemas/mergers/foundations/enriched_wall.py b/generalresearch/incite/schemas/mergers/foundations/enriched_wall.py index 1443f28..71d0eab 100644 --- a/generalresearch/incite/schemas/mergers/foundations/enriched_wall.py +++ b/generalresearch/incite/schemas/mergers/foundations/enriched_wall.py @@ -5,7 +5,7 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.incite.schemas import ARCHIVE_AFTER, PARTITION_ON from generalresearch.locales import Localelator -from generalresearch.models import DeviceType, Source +from generalresearch.models.definitions import DeviceType, Source from generalresearch.models.thl.definitions import ( ReportValue, Status, diff --git a/generalresearch/incite/schemas/mergers/ym_wall_summary.py b/generalresearch/incite/schemas/mergers/ym_wall_summary.py index 16cfc2f..737b925 100644 --- a/generalresearch/incite/schemas/mergers/ym_wall_summary.py +++ b/generalresearch/incite/schemas/mergers/ym_wall_summary.py @@ -6,7 +6,7 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.incite.schemas import ARCHIVE_AFTER from generalresearch.locales import Localelator -from generalresearch.models import Source +from generalresearch.models.definitions import Source COUNTRY_ISOS: set[str] = Localelator().get_all_countries() kosovo = "xk" diff --git a/generalresearch/incite/schemas/thl_web.py b/generalresearch/incite/schemas/thl_web.py index 30c7076..36ee8e9 100644 --- a/generalresearch/incite/schemas/thl_web.py +++ b/generalresearch/incite/schemas/thl_web.py @@ -6,7 +6,7 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index, MultiIndex from generalresearch.incite.schemas import ARCHIVE_AFTER, ORDER_KEY from generalresearch.locales import Localelator -from generalresearch.models import DeviceType, Source +from generalresearch.models.definitions import DeviceType, Source from generalresearch.models.thl.definitions import ( ReportValue, SessionAdjustedStatus, diff --git a/generalresearch/managers/cint/user_pid.py b/generalresearch/managers/cint/user_pid.py index 4f749a0..0265823 100644 --- a/generalresearch/managers/cint/user_pid.py +++ b/generalresearch/managers/cint/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class CintUserPidManager(UserPidManager): diff --git a/generalresearch/managers/dynata/user_pid.py b/generalresearch/managers/dynata/user_pid.py index aefed34..67ff968 100644 --- a/generalresearch/managers/dynata/user_pid.py +++ b/generalresearch/managers/dynata/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class DynataUserPidManager(UserPidManager): diff --git a/generalresearch/managers/events.py b/generalresearch/managers/events.py index c43a020..30cec0c 100644 --- a/generalresearch/managers/events.py +++ b/generalresearch/managers/events.py @@ -12,7 +12,7 @@ from redis.client import PubSub, Redis from generalresearch.incite.base import LOG from generalresearch.managers.base import RedisManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.events import ( AggregateBySource, EventEnvelope, diff --git a/generalresearch/managers/innovate/user_pid.py b/generalresearch/managers/innovate/user_pid.py index 100b0ca..7544c89 100644 --- a/generalresearch/managers/innovate/user_pid.py +++ b/generalresearch/managers/innovate/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class InnovateUserPidManager(UserPidManager): diff --git a/generalresearch/managers/marketplace/user_pid.py b/generalresearch/managers/marketplace/user_pid.py index fe24d38..00dae8a 100644 --- a/generalresearch/managers/marketplace/user_pid.py +++ b/generalresearch/managers/marketplace/user_pid.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING from uuid import UUID from generalresearch.managers.base import SqlManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source if TYPE_CHECKING: from generalresearch.sql_helper import SqlHelper diff --git a/generalresearch/managers/morning/user_pid.py b/generalresearch/managers/morning/user_pid.py index 78de3bd..5896734 100644 --- a/generalresearch/managers/morning/user_pid.py +++ b/generalresearch/managers/morning/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class MorningUserPidManager(UserPidManager): diff --git a/generalresearch/managers/network/label.py b/generalresearch/managers/network/label.py index cec59ad..aed5ff6 100644 --- a/generalresearch/managers/network/label.py +++ b/generalresearch/managers/network/label.py @@ -5,11 +5,12 @@ from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING from psycopg import sql -from pydantic import IPvAnyNetwork, TypeAdapter +from pydantic import TypeAdapter from generalresearch.managers.base import PostgresManager from generalresearch.models.custom_types import ( - AwareDatetimeISO, + IPvAnyAddressStr, + IPvAnyNetwork, IPvAnyNetworkStr, ) from generalresearch.models.network.label import IPLabel @@ -17,7 +18,6 @@ from generalresearch.models.network.label import IPLabel if TYPE_CHECKING: from generalresearch.models.custom_types import ( AwareDatetimeISO, - IPvAnyNetworkStr, ) from generalresearch.models.network.label import IPLabelKind, IPLabelSource diff --git a/generalresearch/managers/precision/user_pid.py b/generalresearch/managers/precision/user_pid.py index 50e97e6..ed2d58d 100644 --- a/generalresearch/managers/precision/user_pid.py +++ b/generalresearch/managers/precision/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class PrecisionUserPidManager(UserPidManager): diff --git a/generalresearch/managers/prodege/user_pid.py b/generalresearch/managers/prodege/user_pid.py index 7c92e28..c18c109 100644 --- a/generalresearch/managers/prodege/user_pid.py +++ b/generalresearch/managers/prodege/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class ProdegeUserPidManager(UserPidManager): diff --git a/generalresearch/managers/repdata/user_pid.py b/generalresearch/managers/repdata/user_pid.py index 9d53897..5fdeccf 100644 --- a/generalresearch/managers/repdata/user_pid.py +++ b/generalresearch/managers/repdata/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class RepdataUserPidManager(UserPidManager): diff --git a/generalresearch/managers/sago/user_pid.py b/generalresearch/managers/sago/user_pid.py index 311abb7..b7ce771 100644 --- a/generalresearch/managers/sago/user_pid.py +++ b/generalresearch/managers/sago/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class SagoUserPidManager(UserPidManager): diff --git a/generalresearch/managers/spectrum/user_pid.py b/generalresearch/managers/spectrum/user_pid.py index 495e73c..980c28d 100644 --- a/generalresearch/managers/spectrum/user_pid.py +++ b/generalresearch/managers/spectrum/user_pid.py @@ -1,5 +1,5 @@ from generalresearch.managers.marketplace.user_pid import UserPidManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source class SpectrumUserPidManager(UserPidManager): diff --git a/generalresearch/managers/thl/buyer.py b/generalresearch/managers/thl/buyer.py index 1e20e2f..38214c6 100644 --- a/generalresearch/managers/thl/buyer.py +++ b/generalresearch/managers/thl/buyer.py @@ -8,7 +8,7 @@ from generalresearch.managers.base import Permission, PostgresManager from generalresearch.models.thl.survey.buyer import Buyer if TYPE_CHECKING: - from generalresearch.models import Source + from generalresearch.models.definitions import Source from generalresearch.pg_helper import PostgresConfig diff --git a/generalresearch/managers/thl/cashout_method.py b/generalresearch/managers/thl/cashout_method.py index e701da3..c12c920 100644 --- a/generalresearch/managers/thl/cashout_method.py +++ b/generalresearch/managers/thl/cashout_method.py @@ -9,10 +9,10 @@ from uuid import UUID, uuid4 from pydantic import NonNegativeInt from generalresearch.managers.base import PostgresManager -from generalresearch.models.thl.wallet import PayoutType from generalresearch.models.thl.wallet.cashout_method import ( CashoutMethod, ) +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.models.thl.user import User diff --git a/generalresearch/managers/thl/contest_manager.py b/generalresearch/managers/thl/contest_manager.py index 3f85d31..68b2cf0 100644 --- a/generalresearch/managers/thl/contest_manager.py +++ b/generalresearch/managers/thl/contest_manager.py @@ -14,7 +14,9 @@ from generalresearch.models.thl.contest import ( ContestPrize, ContestWinner, ) +from generalresearch.models.thl.contest.contest_entry import ContestEntry from generalresearch.models.thl.contest.definitions import ( + ContestEntryType, ContestStatus, ContestType, ) @@ -35,8 +37,6 @@ from generalresearch.models.thl.contest.milestone import ( MilestoneUserView, ) from generalresearch.models.thl.contest.raffle import ( - ContestEntry, - ContestEntryType, RaffleContest, RaffleUserView, ) diff --git a/generalresearch/managers/thl/ledger_manager/ledger.py b/generalresearch/managers/thl/ledger_manager/ledger.py index 6cb4b28..f2455d4 100644 --- a/generalresearch/managers/thl/ledger_manager/ledger.py +++ b/generalresearch/managers/thl/ledger_manager/ledger.py @@ -13,7 +13,6 @@ from pydantic import AwareDatetime, NonNegativeInt, PositiveInt from redis.exceptions import LockError, LockNotOwnedError from generalresearch.currency import LedgerCurrency -from generalresearch.managers import parse_order_by from generalresearch.managers.base import ( Permission, PostgresManager, @@ -28,6 +27,7 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionFlagAlreadyExistsError, LedgerTransactionReleaseLockError, ) +from generalresearch.managers.utils import parse_order_by from generalresearch.models.custom_types import check_valid_uuid from generalresearch.models.thl.ledger import ( LedgerAccount, diff --git a/generalresearch/managers/thl/ledger_manager/thl_ledger.py b/generalresearch/managers/thl/ledger_manager/thl_ledger.py index 7aed619..bd27acf 100644 --- a/generalresearch/managers/thl/ledger_manager/thl_ledger.py +++ b/generalresearch/managers/thl/ledger_manager/thl_ledger.py @@ -52,7 +52,7 @@ from generalresearch.models.thl.ledger import ( ) from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import Status -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/managers/thl/payout.py b/generalresearch/managers/thl/payout.py index 2914ba4..1749783 100644 --- a/generalresearch/managers/thl/payout.py +++ b/generalresearch/managers/thl/payout.py @@ -31,11 +31,11 @@ from generalresearch.models.thl.payout import ( PayoutEvent, UserPayoutEvent, ) -from generalresearch.models.thl.wallet import PayoutType from generalresearch.models.thl.wallet.cashout_method import ( CashMailOrderData, CashoutRequestInfo, ) +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.managers.thl.ledger_manager.thl_ledger import ( diff --git a/generalresearch/managers/thl/session.py b/generalresearch/managers/thl/session.py index 7f17252..41d3893 100644 --- a/generalresearch/managers/thl/session.py +++ b/generalresearch/managers/thl/session.py @@ -10,12 +10,12 @@ from faker import Faker from psycopg import sql from pydantic import NonNegativeInt, PositiveInt -from generalresearch.managers import parse_order_by from generalresearch.managers.base import ( Permission, PostgresManager, ) from generalresearch.managers.thl.product import ProductManager +from generalresearch.managers.utils import parse_order_by from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.session import ( Session, @@ -28,8 +28,8 @@ from generalresearch.models.thl.task_status import ( from generalresearch.models.thl.user import User if TYPE_CHECKING: - from generalresearch.models import DeviceType from generalresearch.models.custom_types import UUIDStr + from generalresearch.models.definitions import DeviceType from generalresearch.models.thl.definitions import ( SessionStatusCode2, Status, diff --git a/generalresearch/managers/thl/survey.py b/generalresearch/managers/thl/survey.py index eacb345..92777e5 100644 --- a/generalresearch/managers/thl/survey.py +++ b/generalresearch/managers/thl/survey.py @@ -13,7 +13,7 @@ from pydantic import NonNegativeInt from generalresearch.managers.base import Permission, PostgresManager from generalresearch.managers.thl.buyer import BuyerManager from generalresearch.managers.thl.category import CategoryManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.survey.model import ( Survey, SurveyStat, diff --git a/generalresearch/managers/thl/task_adjustment.py b/generalresearch/managers/thl/task_adjustment.py index e3f382d..d0d83cb 100644 --- a/generalresearch/managers/thl/task_adjustment.py +++ b/generalresearch/managers/thl/task_adjustment.py @@ -6,12 +6,12 @@ from decimal import Decimal from functools import cached_property from typing import TYPE_CHECKING -from generalresearch.managers import parse_order_by from generalresearch.managers.base import ( PostgresManager, ) from generalresearch.managers.thl.session import SessionManager from generalresearch.managers.thl.wall import WallManager +from generalresearch.managers.utils import parse_order_by from generalresearch.models.thl.definitions import ( Status, WallAdjustedStatus, diff --git a/generalresearch/managers/thl/wall.py b/generalresearch/managers/thl/wall.py index b9dc94d..83697f5 100644 --- a/generalresearch/managers/thl/wall.py +++ b/generalresearch/managers/thl/wall.py @@ -14,12 +14,12 @@ from psycopg import sql from psycopg.rows import dict_row from pydantic import AwareDatetime, PositiveInt -from generalresearch.managers import parse_order_by from generalresearch.managers.base import ( PostgresManager, PostgresManagerWithRedis, ) -from generalresearch.models import Source +from generalresearch.managers.utils import parse_order_by +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( WallAdjustedStatus, ) diff --git a/generalresearch/managers/thl/wallet/__init__.py b/generalresearch/managers/thl/wallet/__init__.py index 457483f..f805872 100644 --- a/generalresearch/managers/thl/wallet/__init__.py +++ b/generalresearch/managers/thl/wallet/__init__.py @@ -6,7 +6,7 @@ from generalresearch.managers.thl.wallet.approve import ( approve_paypal_order, ) from generalresearch.models.thl.definitions import PayoutStatus -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.managers.thl.ledger_manager.thl_ledger import ( diff --git a/generalresearch/managers/utils.py b/generalresearch/managers/utils.py new file mode 100644 index 0000000..bc745fd --- /dev/null +++ b/generalresearch/managers/utils.py @@ -0,0 +1,16 @@ +def parse_order_by(order_by_str: str) -> str: + """ + Converts django-rest-framework ordering str to mysql clause + :param order_by_str: e.g. 'created,-name' + :return: mysql clause e.g. ORDER BY created ASC, name DESC + """ + fields = order_by_str.split(",") + + order_clause = [] + for field in fields: + if field.startswith("-"): + order_clause.append(f"{field[1:]} DESC") + else: + order_clause.append(f"{field} ASC") + + return "ORDER BY " + ", ".join(order_clause) diff --git a/generalresearch/models/cint/question.py b/generalresearch/models/cint/question.py index 44efd13..ab46653 100644 --- a/generalresearch/models/cint/question.py +++ b/generalresearch/models/cint/question.py @@ -8,7 +8,7 @@ from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator -from generalresearch.models import Source, string_utils +from generalresearch.models.definitions import Source, string_utils from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, diff --git a/generalresearch/models/cint/survey.py b/generalresearch/models/cint/survey.py index 8c8f882..ebba09e 100644 --- a/generalresearch/models/cint/survey.py +++ b/generalresearch/models/cint/survey.py @@ -18,7 +18,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import Source, TaskCalculationType +from generalresearch.models.definitions import Source, TaskCalculationType from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask from generalresearch.models.thl.survey.condition import ( diff --git a/generalresearch/models/custom_types.py b/generalresearch/models/custom_types.py index 5e4db3e..680a99c 100644 --- a/generalresearch/models/custom_types.py +++ b/generalresearch/models/custom_types.py @@ -20,7 +20,7 @@ from pydantic.functional_validators import AfterValidator, BeforeValidator from pydantic.networks import IPvAnyNetwork, UrlConstraints from pydantic_core import MultiHostHost, Url -from generalresearch.models import DeviceType, Source +from generalresearch.models.definitions import DeviceType, Source HOSTNAME_REGEX = re.compile( r"^[a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(\.[a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)*$" diff --git a/generalresearch/models/definitions.py b/generalresearch/models/definitions.py new file mode 100644 index 0000000..c0348d7 --- /dev/null +++ b/generalresearch/models/definitions.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from enum import IntEnum, StrEnum + +from generalresearch.utils.enum import ReprEnumMeta + + +class Source(StrEnum, metaclass=ReprEnumMeta): + # The external marketplace, or the source of the survey / work. + # Max length of the value is 2. + GRS = "g" + CINT = "c" + DALIA = "a" # deprecated + DYNATA = "d" + ETX = "et" + FULL_CIRCLE = "f" + INNOVATE = "i" + LUCID = "l" + MORNING_CONSULT = "m" + OPEN_LABS = "n" + POLLFISH = "o" + PRECISION = "e" + PRODEGE_USER = "r" # deprecated + PRODEGE = "pr" # using 'r' for vendor_wall + PULLEY = "p" # deprecated + REPDATA = "rd" # using 'q' for vendor_wall + SAGO = "h" + SPECTRUM = "s" + TESTING = "t" # Used internally for testing + TESTING2 = "u" # Used internally for testing + WXET = "w" + + +class DebitKey(IntEnum, metaclass=ReprEnumMeta): + # The debit key for marketplaces + CINT = 8 + DALIA = 9 + DYNATA = 6 + # ETX = None + FULL_CIRCLE = 15 + INNOVATE = 7 + LUCID = 0 + MORNING_CONSULT = 12 + # OPEN_LABS = None + POLLFISH = 13 + PRECISION = 14 + PRODEGE = 11 + SAGO = 10 + SPECTRUM = 5 + # WXET = None + + +class DeviceType(IntEnum, metaclass=ReprEnumMeta): + UNKNOWN = 0 + MOBILE = 1 + DESKTOP = 2 + TABLET = 3 + + +class LogicalOperator(StrEnum, metaclass=ReprEnumMeta): + OR = "OR" + AND = "AND" + # There is currently no use case for NOT. See MarketplaceCondition.explain_not + NOT = "NOT" + + +class TaskStatus(StrEnum, metaclass=ReprEnumMeta): + # A survey is live if it is open and, given all conditions are met, is + # possible to send in traffic. All other statuses are just variants of + # NOT Live (not accepting traffic) + LIVE = "LIVE" + + # This is a generic NOT Live status. A marketplace may use other more + # specific statuses but in practice they don't matter because all we care + # about is if the task is LIVE. + NOT_LIVE = "NOT_LIVE" + + # We need a status to mark if a survey we thought was live does not come + # back from the API, we'll mark it as NOT_FOUND. + NOT_FOUND = "NOT_FOUND" + + +class TaskCalculationType(StrEnum): + COMPLETES = "COMPLETES" + STARTS = "STARTS" + + @classmethod + def from_api(cls, v: str) -> TaskCalculationType: + return { + "complete": cls.COMPLETES, + "completes": cls.COMPLETES, + "survey start": cls.STARTS, + "survey starts": cls.STARTS, + "start": cls.STARTS, + "prescreens": cls.STARTS, + "prescreen": cls.STARTS, + }[v.lower()] + + @classmethod + def prodege_from_api(cls, v: int) -> TaskCalculationType: + return {1: cls.COMPLETES, 2: cls.STARTS}[v] + + @classmethod + def innovate_from_api(cls, v: int) -> TaskCalculationType: + return {0: cls.COMPLETES, 1: cls.STARTS}[v] + + +class URLQueryKey(StrEnum, metaclass=ReprEnumMeta): + PRODUCT_ID = "39057c8b" + PRODUCT_USER_ID = "c184efc0" + SESSION_ID = "0bb50182" + + +MAX_INT32 = 2**31 diff --git a/generalresearch/models/device.py b/generalresearch/models/device.py index cc15eee..432c897 100644 --- a/generalresearch/models/device.py +++ b/generalresearch/models/device.py @@ -1,6 +1,6 @@ from user_agents import parse as parse_ua -from generalresearch.models import DeviceType +from generalresearch.models.definitions import DeviceType def parse_device_from_useragent(user_agent: str) -> DeviceType: diff --git a/generalresearch/models/dynata/question.py b/generalresearch/models/dynata/question.py index 1ed560a..60c7366 100644 --- a/generalresearch/models/dynata/question.py +++ b/generalresearch/models/dynata/question.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, Field, PositiveInt, field_validator, model_validator -from generalresearch.models import MAX_INT32, Source +from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion if TYPE_CHECKING: diff --git a/generalresearch/models/dynata/survey.py b/generalresearch/models/dynata/survey.py index 70e3659..4174d31 100644 --- a/generalresearch/models/dynata/survey.py +++ b/generalresearch/models/dynata/survey.py @@ -19,7 +19,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.dynata import DynataStatus from generalresearch.models.thl.demographics import ( Gender, @@ -31,7 +31,6 @@ from generalresearch.models.thl.survey.condition import ( ) if TYPE_CHECKING: - from generalresearch.models import TaskCalculationType from generalresearch.models.custom_types import ( AlphaNumStr, AlphaNumStrSet, @@ -39,6 +38,7 @@ if TYPE_CHECKING: CoercedStr, DeviceTypes, ) + from generalresearch.models.definitions import TaskCalculationType logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/dynata/task_collection.py b/generalresearch/models/dynata/task_collection.py index 94868bb..c6cdc19 100644 --- a/generalresearch/models/dynata/task_collection.py +++ b/generalresearch/models/dynata/task_collection.py @@ -6,7 +6,7 @@ import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator -from generalresearch.models import TaskCalculationType +from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.dynata import DynataStatus from generalresearch.models.thl.survey.task_collection import ( TaskCollection, diff --git a/generalresearch/models/events.py b/generalresearch/models/events.py index 8d059f9..34f6be8 100644 --- a/generalresearch/models/events.py +++ b/generalresearch/models/events.py @@ -14,12 +14,12 @@ from pydantic import ( ) if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import ( AwareDatetimeISO, CountryISOLike, UUIDStr, ) + from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( SessionStatusCode2, Status, diff --git a/generalresearch/models/gr/business.py b/generalresearch/models/gr/business.py index 146d690..e11c54d 100644 --- a/generalresearch/models/gr/business.py +++ b/generalresearch/models/gr/business.py @@ -11,7 +11,7 @@ from uuid import uuid4 import pandas as pd import pyarrow as pa -from dask.distributed import Client +from dask.distributed import Client as DaskClient from psycopg.cursor import Cursor from psycopg.rows import dict_row from pydantic import BaseModel, ConfigDict, Field, PositiveInt, ValidationError @@ -24,6 +24,11 @@ from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) from generalresearch.models.admin.request import ReportRequest, ReportType +from generalresearch.models.custom_types import ( + AwareDatetime, + UUIDStr, + UUIDStrCoerce, +) from generalresearch.models.gr.team import Team from generalresearch.models.thl.finance import BusinessBalances, POPFinancial from generalresearch.models.thl.ledger import OrderBy @@ -32,11 +37,6 @@ from generalresearch.utils.enum import ReprEnumMeta if TYPE_CHECKING: from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge - from generalresearch.models.custom_types import ( - AwareDatetime, - UUIDStr, - UUIDStrCoerce, - ) from generalresearch.models.thl.ledger import LedgerAccount from generalresearch.models.thl.payout import BusinessPayoutEvent from generalresearch.pg_helper import PostgresConfig @@ -354,7 +354,7 @@ class Business(BaseModel): thl_pg_config: PostgresConfig, lm: LedgerManager, ds: GRLDatasets, - client: Client, + client: DaskClient, pop_ledger: PopLedgerMerge | None = None, at_timestamp: AwareDatetime | None = None, ) -> None: @@ -464,7 +464,7 @@ class Business(BaseModel): thl_pg_config: PostgresConfig, thl_lm: ThlLedgerManager, ds: GRLDatasets, - client: Client, + client: DaskClient, pop_ledger: PopLedgerMerge | None = None, ) -> None: """This is very similar to the Product POP Financial endpoint; however, @@ -518,7 +518,7 @@ class Business(BaseModel): self, thl_pg_config: PostgresConfig, ds: GRLDatasets, - client: Client, + client: DaskClient, mnt_gr_api: Path, enriched_session: EnrichedSessionMerge | None = None, ) -> None: @@ -561,7 +561,7 @@ class Business(BaseModel): self, thl_pg_config: PostgresConfig, ds: GRLDatasets, - client: Client, + client: DaskClient, mnt_gr_api: Path, enriched_wall: EnrichedWallMerge | None = None, ) -> None: @@ -633,7 +633,7 @@ class Business(BaseModel): pg_config: PostgresConfig, thl_web_rr: PostgresConfig, redis_config: RedisConfig, - client: Client, + client: DaskClient, ds: GRLDatasets, lm: LedgerManager, thl_lm: ThlLedgerManager, diff --git a/generalresearch/models/innovate/question.py b/generalresearch/models/innovate/question.py index fc89524..6423399 100644 --- a/generalresearch/models/innovate/question.py +++ b/generalresearch/models/innovate/question.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, diff --git a/generalresearch/models/innovate/survey.py b/generalresearch/models/innovate/survey.py index 60921df..e718dda 100644 --- a/generalresearch/models/innovate/survey.py +++ b/generalresearch/models/innovate/survey.py @@ -24,7 +24,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import ( +from generalresearch.models.definitions import ( LogicalOperator, Source, ) @@ -41,15 +41,15 @@ from generalresearch.models.thl.survey.condition import ( ) if TYPE_CHECKING: - from generalresearch.models import ( - TaskCalculationType, - ) from generalresearch.models.custom_types import ( AlphaNumStrSet, AwareDatetimeISO, CoercedStr, DeviceTypes, ) + from generalresearch.models.definitions import ( + TaskCalculationType, + ) from generalresearch.models.innovate.question import InnovateQuestionID logging.basicConfig() diff --git a/generalresearch/models/legacy/bucket.py b/generalresearch/models/legacy/bucket.py index f20a769..5f53b89 100644 --- a/generalresearch/models/legacy/bucket.py +++ b/generalresearch/models/legacy/bucket.py @@ -15,7 +15,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.stats import StatisticalSummary if TYPE_CHECKING: diff --git a/generalresearch/models/legacy/questions.py b/generalresearch/models/legacy/questions.py index bebd28f..c333804 100644 --- a/generalresearch/models/legacy/questions.py +++ b/generalresearch/models/legacy/questions.py @@ -219,7 +219,7 @@ class UserQuestionAnswers(BaseModel): self.user = res def prefetch_wall(self, wm: WallManager) -> None: - from generalresearch.models import Source + from generalresearch.models.definitions import Source res: Wall | None = wm.get_from_uuid_if_exists(wall_uuid=self.session_id) diff --git a/generalresearch/models/lucid/question.py b/generalresearch/models/lucid/question.py index 98f535b..c1b9e52 100644 --- a/generalresearch/models/lucid/question.py +++ b/generalresearch/models/lucid/question.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Literal, Self from pydantic import BaseModel, Field, field_validator, model_validator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, ) diff --git a/generalresearch/models/lucid/survey.py b/generalresearch/models/lucid/survey.py index 0f03e31..a04e529 100644 --- a/generalresearch/models/lucid/survey.py +++ b/generalresearch/models/lucid/survey.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Self from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, diff --git a/generalresearch/models/morning/question.py b/generalresearch/models/morning/question.py index 748fcc6..909992f 100644 --- a/generalresearch/models/morning/question.py +++ b/generalresearch/models/morning/question.py @@ -6,7 +6,7 @@ from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator from generalresearch.locales import Localelator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, diff --git a/generalresearch/models/morning/survey.py b/generalresearch/models/morning/survey.py index 1e217f6..25accb6 100644 --- a/generalresearch/models/morning/survey.py +++ b/generalresearch/models/morning/survey.py @@ -25,7 +25,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.morning import MorningStatus from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask diff --git a/generalresearch/models/pollfish/question.py b/generalresearch/models/pollfish/question.py index 3b658fd..f0c733c 100644 --- a/generalresearch/models/pollfish/question.py +++ b/generalresearch/models/pollfish/question.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Literal, Self from pydantic import BaseModel, Field, model_validator -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion if TYPE_CHECKING: diff --git a/generalresearch/models/precision/question.py b/generalresearch/models/precision/question.py index ba17361..6ed6bbd 100644 --- a/generalresearch/models/precision/question.py +++ b/generalresearch/models/precision/question.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator -from generalresearch.models import Source, string_utils +from generalresearch.models.definitions import Source, string_utils from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, diff --git a/generalresearch/models/precision/survey.py b/generalresearch/models/precision/survey.py index fa30882..a9e34e6 100644 --- a/generalresearch/models/precision/survey.py +++ b/generalresearch/models/precision/survey.py @@ -15,7 +15,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.precision import PrecisionStatus from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask diff --git a/generalresearch/models/prodege/question.py b/generalresearch/models/prodege/question.py index c43b51a..b963785 100644 --- a/generalresearch/models/prodege/question.py +++ b/generalresearch/models/prodege/question.py @@ -18,7 +18,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import MAX_INT32, Source +from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion if TYPE_CHECKING: diff --git a/generalresearch/models/prodege/survey.py b/generalresearch/models/prodege/survey.py index 7e56a9c..e3c765e 100644 --- a/generalresearch/models/prodege/survey.py +++ b/generalresearch/models/prodege/survey.py @@ -20,7 +20,11 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import LogicalOperator, Source, TaskCalculationType +from generalresearch.models.definitions import ( + LogicalOperator, + Source, + TaskCalculationType, +) from generalresearch.models.prodege import ( ProdegePastParticipationType, ProdegeStatus, diff --git a/generalresearch/models/repdata/question.py b/generalresearch/models/repdata/question.py index 8cb1fa7..a578741 100644 --- a/generalresearch/models/repdata/question.py +++ b/generalresearch/models/repdata/question.py @@ -17,7 +17,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import MAX_INT32, Source +from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion if TYPE_CHECKING: diff --git a/generalresearch/models/repdata/survey.py b/generalresearch/models/repdata/survey.py index cea61ed..fc1b649 100644 --- a/generalresearch/models/repdata/survey.py +++ b/generalresearch/models/repdata/survey.py @@ -21,7 +21,7 @@ from pydantic import ( from generalresearch.grpc import timestamp_from_datetime from generalresearch.locales import Localelator -from generalresearch.models import ( +from generalresearch.models.definitions import ( DeviceType, LogicalOperator, Source, diff --git a/generalresearch/models/repdata/task_collection.py b/generalresearch/models/repdata/task_collection.py index 04d79bd..f2cb63b 100644 --- a/generalresearch/models/repdata/task_collection.py +++ b/generalresearch/models/repdata/task_collection.py @@ -6,7 +6,7 @@ import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator -from generalresearch.models import TaskCalculationType +from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.repdata import RepDataStatus from generalresearch.models.thl.survey.task_collection import ( TaskCollection, diff --git a/generalresearch/models/sago/question.py b/generalresearch/models/sago/question.py index cf9ea19..bb51d31 100644 --- a/generalresearch/models/sago/question.py +++ b/generalresearch/models/sago/question.py @@ -18,7 +18,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import MAX_INT32, Source, string_utils +from generalresearch.models.definitions import MAX_INT32, Source, string_utils from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion if TYPE_CHECKING: diff --git a/generalresearch/models/sago/survey.py b/generalresearch/models/sago/survey.py index c2f886a..c9bf431 100644 --- a/generalresearch/models/sago/survey.py +++ b/generalresearch/models/sago/survey.py @@ -18,7 +18,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator -from generalresearch.models import LogicalOperator, Source +from generalresearch.models.definitions import LogicalOperator, Source from generalresearch.models.sago import SagoStatus from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask diff --git a/generalresearch/models/spectrum/question.py b/generalresearch/models/spectrum/question.py index 89fbeb3..9c9bfa0 100644 --- a/generalresearch/models/spectrum/question.py +++ b/generalresearch/models/spectrum/question.py @@ -18,7 +18,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import MAX_INT32, Source, string_utils +from generalresearch.models.definitions import MAX_INT32, Source, string_utils from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, ) diff --git a/generalresearch/models/spectrum/survey.py b/generalresearch/models/spectrum/survey.py index 4daa00b..a02c510 100644 --- a/generalresearch/models/spectrum/survey.py +++ b/generalresearch/models/spectrum/survey.py @@ -10,7 +10,7 @@ from more_itertools import flatten from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator from generalresearch.locales import Localelator -from generalresearch.models import Source, TaskCalculationType +from generalresearch.models.definitions import Source, TaskCalculationType from generalresearch.models.spectrum import SpectrumStatus from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask diff --git a/generalresearch/models/spectrum/task_collection.py b/generalresearch/models/spectrum/task_collection.py index d909292..8e49434 100644 --- a/generalresearch/models/spectrum/task_collection.py +++ b/generalresearch/models/spectrum/task_collection.py @@ -6,7 +6,7 @@ import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator -from generalresearch.models import TaskCalculationType +from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.spectrum import SpectrumStatus from generalresearch.models.thl.survey.task_collection import ( TaskCollection, diff --git a/generalresearch/models/thl/__init__.py b/generalresearch/models/thl/__init__.py index cb04b29..7f2b8a9 100644 --- a/generalresearch/models/thl/__init__.py +++ b/generalresearch/models/thl/__init__.py @@ -8,7 +8,7 @@ from decimal import Decimal # BrokerageProductPayoutEvent, # PayoutEvent, # ) -# from generalresearch.models.thl.product import Product +from generalresearch.models.thl.product import Product # _ = ( # Product, @@ -18,7 +18,7 @@ from decimal import Decimal # POPFinancial, # ) -# Product.model_rebuild() +Product.model_rebuild() # PayoutEvent.model_rebuild() # BrokerageProductPayoutEvent.model_rebuild() diff --git a/generalresearch/models/thl/category.py b/generalresearch/models/thl/category.py index 32841a5..ebfc840 100644 --- a/generalresearch/models/thl/category.py +++ b/generalresearch/models/thl/category.py @@ -1,12 +1,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Self +from typing import Any, Self from uuid import uuid4 from pydantic import BaseModel, Field, PositiveInt, model_validator -if TYPE_CHECKING: - from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.custom_types import UUIDStr class Category(BaseModel, frozen=True): diff --git a/generalresearch/models/thl/contest/contest_entry.py b/generalresearch/models/thl/contest/contest_entry.py index a57b2df..17b288b 100644 --- a/generalresearch/models/thl/contest/contest_entry.py +++ b/generalresearch/models/thl/contest/contest_entry.py @@ -12,10 +12,12 @@ from pydantic import ( ) from generalresearch.currency import USDCent +from generalresearch.models.thl.contest.definitions import ( + ContestEntryType, +) if TYPE_CHECKING: from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr - from generalresearch.models.thl.contest.definitions import ContestEntryType from generalresearch.models.thl.user import User @@ -59,10 +61,7 @@ class ContestEntry(BaseModel): @model_validator(mode="before") @classmethod - def validate_amount_type(cls, data: dict) -> dict: - from generalresearch.models.thl.contest.definitions import ( - ContestEntryType, - ) + def validate_amount_type(cls, data: dict[str, Any]) -> dict[str, Any]: amount = data.get("amount") entry_type = data.get("entry_type") @@ -71,6 +70,7 @@ class ContestEntry(BaseModel): assert isinstance(amount, int) and not isinstance( amount, USDCent ), "amount must be int in ContestEntryType.COUNT" + elif entry_type == ContestEntryType.CASH: # This may be coming from the DB, in which case it is an int. data["amount"] = USDCent(data["amount"]) @@ -79,9 +79,6 @@ class ContestEntry(BaseModel): @computed_field() def amount_str(self) -> str: - from generalresearch.models.thl.contest.definitions import ( - ContestEntryType, - ) if self.entry_type == ContestEntryType.COUNT: return str(self.amount) diff --git a/generalresearch/models/thl/contest/raffle.py b/generalresearch/models/thl/contest/raffle.py index b944740..21bc481 100644 --- a/generalresearch/models/thl/contest/raffle.py +++ b/generalresearch/models/thl/contest/raffle.py @@ -26,11 +26,9 @@ from generalresearch.models.thl.contest.contest import ( ContestBase, ContestUserView, ) -from generalresearch.models.thl.contest.contest_entry import ( - ContestEntryType, -) from generalresearch.models.thl.contest.definitions import ( ContestEndReason, + ContestEntryType, ContestPrizeKind, ContestStatus, ContestType, diff --git a/generalresearch/models/thl/finance.py b/generalresearch/models/thl/finance.py index 4b750da..9e7d2c3 100644 --- a/generalresearch/models/thl/finance.py +++ b/generalresearch/models/thl/finance.py @@ -27,8 +27,9 @@ adjustment_example = random.randint(-1_000, 50 * 100) if TYPE_CHECKING: from generalresearch.currency import USDCent + from generalresearch.managers.thl.product import ProductManager from generalresearch.models.thl.ledger import LedgerAccount - from generalresearch.pg_helper import PostgresConfig + from generalresearch.models.thl.product import Product class AdjustmentType(BaseModel): @@ -516,7 +517,7 @@ class ProductBalances(BaseModel): if isinstance(input_data, pd.Series): return ProductBalances.model_validate(input_data.to_dict()) - elif isinstance(input_data, pd.DataFrame): + else: assert isinstance(input_data.index, pd.DatetimeIndex), "Invalid input data" # The pop merge is grouped by 1min intervals. Therefore, if we take @@ -529,9 +530,6 @@ class ProductBalances(BaseModel): pb.last_event = pq_last_event_close.to_pydatetime() return pb - else: - raise NotImplementedError("Can't handle this input") - def __str__(self) -> str: return ( f"Product: {self.product_id or '—'}\n" @@ -834,19 +832,17 @@ class BusinessBalances(BaseModel): def from_pandas( input_data: pd.DataFrame, accounts: list[LedgerAccount], - thl_pg_config: PostgresConfig, + product_manager: ProductManager, ) -> BusinessBalances: LOG.debug(f"BusinessBalances.from_pandas(input_data={input_data.shape})") from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) - from generalresearch.managers.thl.product import ProductManager from generalresearch.models.thl.ledger import ( AccountType, Direction, ) - from generalresearch.models.thl.product import Product # Validate the input accounts assert len(accounts) > 0, "Must provide accounts" @@ -872,8 +868,7 @@ class BusinessBalances(BaseModel): # Sort the ProductBalances so that they're always in a consistent # sorted order. - pm = ProductManager(pg_config=thl_pg_config) - products: list[Product] = pm.get_by_uuids( + products: list[Product] = product_manager.get_by_uuids( product_uuids=[pb.product_id for pb in product_balances] ) sorted_products_uuids = [ diff --git a/generalresearch/models/thl/ledger.py b/generalresearch/models/thl/ledger.py index c38e83b..fbfb6bb 100644 --- a/generalresearch/models/thl/ledger.py +++ b/generalresearch/models/thl/ledger.py @@ -354,7 +354,7 @@ class LedgerTransaction(BaseModel): def to_user_tx( self, user_account: LedgerAccount, product_id: str, payout_format: str ): - from generalresearch.models.thl.wallet import PayoutType + from generalresearch.models.thl.wallet.definitions import PayoutType d = self.model_dump(include={"created"}) d["tx_type"] = self.metadata.get("tx_type") diff --git a/generalresearch/models/thl/offerwall/__init__.py b/generalresearch/models/thl/offerwall/__init__.py index 0c3d51d..599cc1d 100644 --- a/generalresearch/models/thl/offerwall/__init__.py +++ b/generalresearch/models/thl/offerwall/__init__.py @@ -14,8 +14,8 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source from generalresearch.models.custom_types import IPvAnyAddressStr +from generalresearch.models.definitions import Source from generalresearch.models.thl.locales import ( CountryISO, LanguageISO, diff --git a/generalresearch/models/thl/offerwall/base.py b/generalresearch/models/thl/offerwall/base.py index 1d41ef2..fb0bc77 100644 --- a/generalresearch/models/thl/offerwall/base.py +++ b/generalresearch/models/thl/offerwall/base.py @@ -19,7 +19,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.legacy.bucket import ( Bucket as LegacyBucket, ) diff --git a/generalresearch/models/thl/offerwall/cache.py b/generalresearch/models/thl/offerwall/cache.py index aa18014..2a733c9 100644 --- a/generalresearch/models/thl/offerwall/cache.py +++ b/generalresearch/models/thl/offerwall/cache.py @@ -6,8 +6,8 @@ from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr + from generalresearch.models.definitions import Source from generalresearch.models.thl.offerwall import OfferWallRequest from generalresearch.models.thl.offerwall.base import ( OfferwallBase, diff --git a/generalresearch/models/thl/payout.py b/generalresearch/models/thl/payout.py index 9902af3..128723b 100644 --- a/generalresearch/models/thl/payout.py +++ b/generalresearch/models/thl/payout.py @@ -18,7 +18,7 @@ from pydantic.json_schema import SkipJsonSchema from generalresearch.currency import USDCent from generalresearch.models.thl.definitions import PayoutStatus -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.models.custom_types import ( diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index 988b72d..346a98b 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -38,13 +38,13 @@ from pydantic.json_schema import SkipJsonSchema from generalresearch.currency import USDCent from generalresearch.decorators import LOG -from generalresearch.models import Source from generalresearch.models.custom_types import ( AwareDatetimeISO, CountryISOLike, HttpsUrlStr, UUIDStr, ) +from generalresearch.models.definitions import Source from generalresearch.models.thl.finance import ( POPFinancial, ProductBalances, @@ -63,7 +63,7 @@ from generalresearch.models.thl.payout_format import ( examples as payout_format_examples, ) from generalresearch.models.thl.supplier_tag import SupplierTag -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType from generalresearch.models.utils import decimal_to_usd_cents from generalresearch.redis_helper import RedisConfig diff --git a/generalresearch/models/thl/profiling/marketplace.py b/generalresearch/models/thl/profiling/marketplace.py index 0c1e39b..23501e3 100644 --- a/generalresearch/models/thl/profiling/marketplace.py +++ b/generalresearch/models/thl/profiling/marketplace.py @@ -7,16 +7,16 @@ from typing import TYPE_CHECKING, Any from pydantic import BaseModel, ConfigDict, Field, PositiveInt, computed_field -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import ( AwareDatetimeISO, CountryISOLike, LanguageISOLike, UUIDStr, ) + from generalresearch.models.definitions import Source from generalresearch.models.thl.locales import CountryISO, LanguageISO diff --git a/generalresearch/models/thl/profiling/upk_question.py b/generalresearch/models/thl/profiling/upk_question.py index a73683c..9c7383a 100644 --- a/generalresearch/models/thl/profiling/upk_question.py +++ b/generalresearch/models/thl/profiling/upk_question.py @@ -5,7 +5,7 @@ import json import re from enum import StrEnum from functools import cached_property -from typing import TYPE_CHECKING, Annotated, Any, Literal +from typing import Annotated, Any, Literal from pydantic import ( BaseModel, @@ -17,12 +17,10 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source +from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.definitions import Source from generalresearch.models.thl.category import Category -if TYPE_CHECKING: - from generalresearch.models.custom_types import UUIDStr - class UPKImportance(BaseModel): task_count: int | None = Field( diff --git a/generalresearch/models/thl/profiling/upk_question_answer.py b/generalresearch/models/thl/profiling/upk_question_answer.py index 41895b1..4d07970 100644 --- a/generalresearch/models/thl/profiling/upk_question_answer.py +++ b/generalresearch/models/thl/profiling/upk_question_answer.py @@ -13,7 +13,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 from generalresearch.models.thl.profiling.upk_property import ( Cardinality, PropertyType, diff --git a/generalresearch/models/thl/profiling/user_info.py b/generalresearch/models/thl/profiling/user_info.py index c82e2d2..40b4b17 100644 --- a/generalresearch/models/thl/profiling/user_info.py +++ b/generalresearch/models/thl/profiling/user_info.py @@ -6,8 +6,8 @@ from pydantic import BaseModel, ConfigDict, Field from pydantic.json_schema import SkipJsonSchema if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import AwareDatetimeISO + from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.user_question_answer import ( MarketplaceResearchProfileQuestion, ) diff --git a/generalresearch/models/thl/profiling/user_question_answer.py b/generalresearch/models/thl/profiling/user_question_answer.py index b1868b3..a7c2194 100644 --- a/generalresearch/models/thl/profiling/user_question_answer.py +++ b/generalresearch/models/thl/profiling/user_question_answer.py @@ -14,12 +14,12 @@ from pydantic import ( model_validator, ) -from generalresearch.models import MAX_INT32 +from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr +from generalresearch.models.definitions import MAX_INT32 +from generalresearch.models.thl.locales import CountryISO, LanguageISO if TYPE_CHECKING: - from generalresearch.models import Source - from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr - from generalresearch.models.thl.locales import CountryISO, LanguageISO + from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.upk_question import UpkQuestion diff --git a/generalresearch/models/thl/session.py b/generalresearch/models/thl/session.py index 871e5c4..404cff7 100644 --- a/generalresearch/models/thl/session.py +++ b/generalresearch/models/thl/session.py @@ -18,7 +18,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl import ( decimal_to_int_cents, int_cents_to_decimal, @@ -37,13 +37,13 @@ if TYPE_CHECKING: from generalresearch.managers.thl.ledger_manager.thl_ledger import ( ThlLedgerManager, ) - from generalresearch.models import DeviceType from generalresearch.models.custom_types import ( AwareDatetimeISO, EnumNameSerializer, IPvAnyAddressStr, UUIDStr, ) + from generalresearch.models.definitions import DeviceType from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.definitions import ( ReportValue, diff --git a/generalresearch/models/thl/soft_pair.py b/generalresearch/models/thl/soft_pair.py index 313d374..c0bf2dd 100644 --- a/generalresearch/models/thl/soft_pair.py +++ b/generalresearch/models/thl/soft_pair.py @@ -5,7 +5,7 @@ from enum import Enum from typing import TYPE_CHECKING if TYPE_CHECKING: - from generalresearch.models import Source + from generalresearch.models.definitions import Source from generalresearch.models.dynata.survey import DynataCondition from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, diff --git a/generalresearch/models/thl/survey/__init__.py b/generalresearch/models/thl/survey/__init__.py index d0f2b33..76f819e 100644 --- a/generalresearch/models/thl/survey/__init__.py +++ b/generalresearch/models/thl/survey/__init__.py @@ -18,7 +18,7 @@ from generalresearch.models.thl.survey.condition import ( ) if TYPE_CHECKING: - from generalresearch.models import Source + from generalresearch.models.definitions import Source from generalresearch.models.thl.locales import ( CountryISO, CountryISOs, diff --git a/generalresearch/models/thl/survey/buyer.py b/generalresearch/models/thl/survey/buyer.py index 26846d3..ef309d1 100644 --- a/generalresearch/models/thl/survey/buyer.py +++ b/generalresearch/models/thl/survey/buyer.py @@ -16,7 +16,7 @@ from pydantic import ( ) from scipy.stats import beta as beta_dist -from generalresearch.models import Source +from generalresearch.models.definitions import Source if TYPE_CHECKING: from generalresearch.models.custom_types import ( diff --git a/generalresearch/models/thl/survey/condition.py b/generalresearch/models/thl/survey/condition.py index 514ee64..90cf27b 100644 --- a/generalresearch/models/thl/survey/condition.py +++ b/generalresearch/models/thl/survey/condition.py @@ -17,7 +17,7 @@ from pydantic import ( model_validator, ) -from generalresearch.models import LogicalOperator +from generalresearch.models.definitions import LogicalOperator MarketplaceConditionHash = Annotated[ str, StringConstraints(min_length=7, max_length=7, pattern=r"^[a-f0-9]+$") diff --git a/generalresearch/models/thl/survey/model.py b/generalresearch/models/thl/survey/model.py index 9fa3d8e..8986e4d 100644 --- a/generalresearch/models/thl/survey/model.py +++ b/generalresearch/models/thl/survey/model.py @@ -21,7 +21,6 @@ from generalresearch.models.thl.definitions import StatusCode1 from generalresearch.models.thl.pagination import Page if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import ( AwareDatetimeISO, CountryISOLike, @@ -29,6 +28,7 @@ if TYPE_CHECKING: PropertyCode, SurveyKey, ) + from generalresearch.models.definitions import Source from generalresearch.models.thl.category import Category from generalresearch.models.thl.definitions import Status diff --git a/generalresearch/models/thl/survey/penalty.py b/generalresearch/models/thl/survey/penalty.py index 54edb94..25e07cf 100644 --- a/generalresearch/models/thl/survey/penalty.py +++ b/generalresearch/models/thl/survey/penalty.py @@ -7,11 +7,11 @@ from typing import TYPE_CHECKING, Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, TypeAdapter if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import ( AwareDatetimeISO, UUIDStr, ) + from generalresearch.models.definitions import Source class SurveyPenalty(BaseModel, abc.ABC): diff --git a/generalresearch/models/thl/task_adjustment.py b/generalresearch/models/thl/task_adjustment.py index fa5592e..fee2007 100644 --- a/generalresearch/models/thl/task_adjustment.py +++ b/generalresearch/models/thl/task_adjustment.py @@ -7,14 +7,14 @@ from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 from generalresearch.models.thl.definitions import ( WallAdjustedStatus, ) if TYPE_CHECKING: - from generalresearch.models import Source from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr + from generalresearch.models.definitions import Source class TaskAdjustmentEvent(BaseModel): diff --git a/generalresearch/models/thl/user.py b/generalresearch/models/thl/user.py index 302aa72..1f88dc6 100644 --- a/generalresearch/models/thl/user.py +++ b/generalresearch/models/thl/user.py @@ -20,7 +20,7 @@ from pydantic import ( ) from sentry_sdk import set_tag, set_user -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 if TYPE_CHECKING: from generalresearch.managers.thl.ledger_manager.thl_ledger import ( @@ -253,7 +253,7 @@ class User(BaseModel): # # Delete from db.thl-marketplaces # We need DELETE credentials for all these... - # from generalresearch.models import Source + # from generalresearch.models.definitions import Source # mp_db_table = { # Source.SPECTRUM: "`thl-spectrum`.`spectrum_marketresearchprofilequestion`", # Source.INNOVATE: "`thl-innovate`.`innovate_marketresearchprofilequestion`", diff --git a/generalresearch/models/thl/user_profile.py b/generalresearch/models/thl/user_profile.py index 2dc19b7..c47c6f2 100644 --- a/generalresearch/models/thl/user_profile.py +++ b/generalresearch/models/thl/user_profile.py @@ -13,7 +13,7 @@ from pydantic import ( ) from pydantic.json_schema import SkipJsonSchema -from generalresearch.models import MAX_INT32, Source +from generalresearch.models.definitions import MAX_INT32, Source if TYPE_CHECKING: from generalresearch.models.custom_types import UUIDStr diff --git a/generalresearch/models/thl/user_quality_event.py b/generalresearch/models/thl/user_quality_event.py index 5438740..8c2e25f 100644 --- a/generalresearch/models/thl/user_quality_event.py +++ b/generalresearch/models/thl/user_quality_event.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Literal from pydantic import BaseModel, Field, PositiveInt -from generalresearch.models import MAX_INT32, Source +from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.utils.enum import ReprEnumMeta if TYPE_CHECKING: diff --git a/generalresearch/models/thl/user_streak.py b/generalresearch/models/thl/user_streak.py index 6cd853a..4c09d13 100644 --- a/generalresearch/models/thl/user_streak.py +++ b/generalresearch/models/thl/user_streak.py @@ -19,7 +19,7 @@ from pydantic import ( from pydantic.json_schema import SkipJsonSchema from generalresearch.managers.leaderboard import country_timezone -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 if TYPE_CHECKING: from generalresearch.models.thl.locales import CountryISO diff --git a/generalresearch/models/thl/wallet/cashout_method.py b/generalresearch/models/thl/wallet/cashout_method.py index 1db85e8..9383c36 100644 --- a/generalresearch/models/thl/wallet/cashout_method.py +++ b/generalresearch/models/thl/wallet/cashout_method.py @@ -19,7 +19,7 @@ from pydantic import ( from generalresearch.models.legacy.api_status import StatusResponse from generalresearch.models.thl.definitions import PayoutStatus -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType from generalresearch.utils.enum import ReprEnumMeta if TYPE_CHECKING: @@ -31,7 +31,7 @@ if TYPE_CHECKING: ) from generalresearch.models.thl.locales import CountryISO from generalresearch.models.thl.user import BPUIDStr, User - from generalresearch.models.thl.wallet import Currency + from generalresearch.models.thl.wallet.definitions import Currency logger = logging.getLogger() diff --git a/generalresearch/models/thl/wallet/definitions.py b/generalresearch/models/thl/wallet/definitions.py new file mode 100644 index 0000000..2d1eb8d --- /dev/null +++ b/generalresearch/models/thl/wallet/definitions.py @@ -0,0 +1,87 @@ +from enum import StrEnum + +from generalresearch.utils.enum import ReprEnumMeta + + +class PayoutType(StrEnum, metaclass=ReprEnumMeta): + """ + The method in which the requested payout is delivered. + """ + + # The max size of the db field that holds this value is 14, so please + # don't add new values longer than that! + + # User is paid out to their personal PayPal email address + PAYPAL = "PAYPAL" + # User is paid out via a Tango Gift Card + TANGO = "TANGO" + # DWOLLA + DWOLLA = "DWOLLA" + # A payment is made to a bank account using ACH + ACH = "ACH" + # A payment is made to a bank account using ACH + WIRE = "WIRE" + # A payment is made in cash and mailed to the user. + CASH_IN_MAIL = "CASH_IN_MAIL" + # A payment is made as a prize with some monetary value + PRIZE = "PRIZE" + + # This is used to designate either AMT_BONUS or AMT_HIT + AMT = "AMT" + # Amazon Mechanical Turk as a Bonus + AMT_BONUS = "AMT_BONUS" + # Amazon Mechanical Turk for a HIT + AMT_HIT = "AMT_ASSIGNMENT" + AMT_ASSIGNMENT = "AMT_ASSIGNMENT" + + +class Currency(StrEnum): + # United States Dollar + USD = "USD" + # Canadian Dollar + CAD = "CAD" + # British Pound Sterling + GBP = "GBP" + # Euro + EUR = "EUR" + # Indian Rupee + INR = "INR" + # Australian Dollar + AUD = "AUD" + # Polish Zloty + PLN = "PLN" + # Swedish Krona + SEK = "SEK" + # Singapore Dollar + SGD = "SGD" + # Mexican Peso + MXN = "MXN" + + +CURRENCY_FORMATTER = { + "USD": lambda x: f"${x / 100:,.2f}", + "CAD": lambda x: f"${x / 100:,.2f} CAD", + "GBP": lambda x: f"{x / 100:,.2f} £", + "EUR": lambda x: f"€{x / 100:,.2f}", + "INR": lambda x: f"₹{x / 100:,.2f}", + "AUD": lambda x: f"${x / 100:,.2f} AUD", + "PLN": lambda x: f"{x / 100:,.2f} zł", + "SEK": lambda x: f"{x / 100:,.2f} kr", + "SGD": lambda x: f"${x / 100:,.2f} SGD", + "MXN": lambda x: f"${x / 100:,.2f} MXN", +} + +# The max value user can redeem in one go in foreign currencies. should be < $250 +# in order to avoid exchange rate issues +CURRENCY_MAX_VALUE = { + "USD": 250, + "CAD": 200, + "GBP": 100, + "EUR": 100, + "INR": 10000, + "AUD": 200, + "PLN": 500, + "SEK": 1000, + "SGD": 200, + "MXN": 4000, +} diff --git a/generalresearch/models/thl/wallet/payout.py b/generalresearch/models/thl/wallet/payout.py index 7301b31..79c50e1 100644 --- a/generalresearch/models/thl/wallet/payout.py +++ b/generalresearch/models/thl/wallet/payout.py @@ -16,7 +16,7 @@ from pydantic import ( from generalresearch.currency import USDCent from generalresearch.models.thl.definitions import PayoutStatus -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType if TYPE_CHECKING: from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr diff --git a/generalresearch/schemas/survey_stats.py b/generalresearch/schemas/survey_stats.py index b3acf34..dd592d4 100644 --- a/generalresearch/schemas/survey_stats.py +++ b/generalresearch/schemas/survey_stats.py @@ -2,7 +2,7 @@ import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator -from generalresearch.models import Source +from generalresearch.models.definitions import Source COUNTRY_ISOS = Localelator().get_all_countries() kosovo = "xk" diff --git a/generalresearch/wall_status_codes/__init__.py b/generalresearch/wall_status_codes/__init__.py index 37f3960..cca1a19 100644 --- a/generalresearch/wall_status_codes/__init__.py +++ b/generalresearch/wall_status_codes/__init__.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import Status, StatusCode1 from generalresearch.wall_status_codes import ( cint, diff --git a/test_utils/conftest.py b/test_utils/conftest.py index f55fe11..397d98f 100644 --- a/test_utils/conftest.py +++ b/test_utils/conftest.py @@ -9,6 +9,7 @@ from collections.abc import Callable, Generator from datetime import UTC, datetime, timedelta from os.path import join as pjoin from pathlib import Path +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -17,23 +18,13 @@ from dotenv import load_dotenv from pydantic import MariaDBDsn, PostgresDsn, TypeAdapter from pytest import TempPathFactory -from generalresearch.config import GRLBaseSettings from generalresearch.currency import USDCent from generalresearch.models.custom_types import InternalHostname, PostgresDict -from generalresearch.pg_helper import PostgresConfig from generalresearch.sql_helper import SqlHelper -# -- redis notes from jenkins file -# sh "redis-cli -u ${env.THL_REDIS} FLUSHDB" -# sh "redis-cli -u ${env.GR_REDIS} FLUSHDB" - -# script { -# env.GR_REDIS_DB = new Random().nextInt(1024).toString() -# env.GR_REDIS = "redis://${env.REDIS}:6379/${env.GR_REDIS_DB}" -# echo "Using GR Redis: ${env.GR_REDIS}" -# if (sh(script: "redis-cli -u ${env.GR_REDIS} SET jenkins_lock 1 NX EX 3600", returnStdout: true).trim() != 'OK') -# error('Redis already locked... aborting.') -# } +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.pg_helper import PostgresConfig @pytest.fixture(scope="session") diff --git a/test_utils/grliq/conftest.py b/test_utils/grliq/conftest.py index 891b73c..bb1a167 100644 --- a/test_utils/grliq/conftest.py +++ b/test_utils/grliq/conftest.py @@ -2,19 +2,15 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 import pytest from pydantic import PostgresDsn -from generalresearch.config import GRLBaseSettings from generalresearch.grliq.managers.forensic_data import ( GrlIqDataManager, ) -from generalresearch.grliq.managers.forensic_events import ( - GrlIqEventManager, -) from generalresearch.grliq.managers.forensic_results import ( GrlIqCategoryResultsReader, ) @@ -25,6 +21,12 @@ from generalresearch.grliq.models.forensic_result import ( ) from generalresearch.pg_helper import PostgresConfig +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.grliq.managers.forensic_events import ( + GrlIqEventManager, + ) + # === Miscellaneous === diff --git a/test_utils/incite/collections/conftest.py b/test_utils/incite/collections/conftest.py index 631bb7b..f490e14 100644 --- a/test_utils/incite/collections/conftest.py +++ b/test_utils/incite/collections/conftest.py @@ -6,12 +6,11 @@ from typing import TYPE_CHECKING import pytest -from generalresearch.pg_helper import PostgresConfig from test_utils.conftest import clear_directory if TYPE_CHECKING: from generalresearch.incite.base import DFCollectionType, GRLDatasets - from generalresearch.incite.collections import DFCollection + from generalresearch.incite.collections.base import DFCollection from generalresearch.incite.collections.thl_web import ( AuditLogDFCollection, LedgerDFCollection, @@ -20,6 +19,7 @@ if TYPE_CHECKING: UserDFCollection, WallDFCollection, ) + from generalresearch.pg_helper import PostgresConfig @pytest.fixture diff --git a/test_utils/incite/mergers/conftest.py b/test_utils/incite/mergers/conftest.py index 1f88804..4eb3f2d 100644 --- a/test_utils/incite/mergers/conftest.py +++ b/test_utils/incite/mergers/conftest.py @@ -2,37 +2,40 @@ from __future__ import annotations from collections.abc import Callable from datetime import datetime, timedelta +from typing import TYPE_CHECKING import pytest -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.mergers.base import MergeType -from generalresearch.incite.mergers.foundations.enriched_session import ( - EnrichedSessionMerge, -) -from generalresearch.incite.mergers.foundations.enriched_task_adjust import ( - EnrichedTaskAdjustMerge, -) -from generalresearch.incite.mergers.foundations.enriched_wall import ( - EnrichedWallMerge, -) -from generalresearch.incite.mergers.foundations.user_id_product import ( - UserIdProductMerge, -) -from generalresearch.incite.mergers.pop_ledger import ( - PopLedgerMerge, - PopLedgerMergeItem, -) -from generalresearch.incite.mergers.ym_survey_wall import ( - YMSurveyWallMerge, - YMSurveyWallMergeCollectionItem, -) -from generalresearch.incite.mergers.ym_wall_summary import ( - YMWallSummaryMerge, - YMWallSummaryMergeItem, -) from test_utils.conftest import clear_directory +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.mergers.base import MergeType + from generalresearch.incite.mergers.foundations.enriched_session import ( + EnrichedSessionMerge, + ) + from generalresearch.incite.mergers.foundations.enriched_task_adjust import ( + EnrichedTaskAdjustMerge, + ) + from generalresearch.incite.mergers.foundations.enriched_wall import ( + EnrichedWallMerge, + ) + from generalresearch.incite.mergers.foundations.user_id_product import ( + UserIdProductMerge, + ) + from generalresearch.incite.mergers.pop_ledger import ( + PopLedgerMerge, + PopLedgerMergeItem, + ) + from generalresearch.incite.mergers.ym_survey_wall import ( + YMSurveyWallMerge, + YMSurveyWallMergeCollectionItem, + ) + from generalresearch.incite.mergers.ym_wall_summary import ( + YMWallSummaryMerge, + YMWallSummaryMergeItem, + ) + # -------------------------- # Merges # -------------------------- diff --git a/test_utils/managers/cashout_methods.py b/test_utils/managers/cashout_methods.py index 238cdda..adf82f4 100644 --- a/test_utils/managers/cashout_methods.py +++ b/test_utils/managers/cashout_methods.py @@ -6,11 +6,11 @@ from uuid import uuid4 import pytest -from generalresearch.models.thl.wallet import Currency, PayoutType from generalresearch.models.thl.wallet.cashout_method import ( CashoutMethod, TangoCashoutMethodData, ) +from generalresearch.models.thl.wallet.definitions import Currency, PayoutType @pytest.fixture(scope="session") diff --git a/test_utils/managers/conftest.py b/test_utils/managers/conftest.py index 9c6a1a7..ed771c7 100644 --- a/test_utils/managers/conftest.py +++ b/test_utils/managers/conftest.py @@ -1,41 +1,44 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.gr.business import ( - BusinessAddressManager, - BusinessBankAccountManager, - BusinessManager, -) -from generalresearch.managers.gr.team import ( - MembershipManager, - TeamManager, -) -from generalresearch.managers.spectrum.survey import SpectrumSurveyManager -from generalresearch.managers.thl.buyer import BuyerManager from generalresearch.managers.thl.cashout_method import ( CashoutMethodManager, ) -from generalresearch.managers.thl.ipinfo import ( - GeoIpInfoManager, - IPGeonameManager, - IPInformationManager, -) from generalresearch.managers.thl.user_streak import ( UserStreakManager, ) -from generalresearch.managers.thl.userhealth import ( - AuditLogManager, - IPRecordManager, - UserIpHistoryManager, -) -from generalresearch.models import Source -from generalresearch.models.thl.wallet.cashout_method import CashoutMethod -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig -from generalresearch.sql_helper import SqlHelper +from generalresearch.models.definitions import Source + +if TYPE_CHECKING: + from generalresearch.managers.gr.business import ( + BusinessAddressManager, + BusinessBankAccountManager, + BusinessManager, + ) + from generalresearch.managers.gr.team import ( + MembershipManager, + TeamManager, + ) + from generalresearch.managers.spectrum.survey import SpectrumSurveyManager + from generalresearch.managers.thl.buyer import BuyerManager + from generalresearch.managers.thl.ipinfo import ( + GeoIpInfoManager, + IPGeonameManager, + IPInformationManager, + ) + from generalresearch.managers.thl.userhealth import ( + AuditLogManager, + IPRecordManager, + UserIpHistoryManager, + ) + from generalresearch.models.thl.wallet.cashout_method import CashoutMethod + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig + from generalresearch.sql_helper import SqlHelper # === THL === diff --git a/test_utils/managers/contest/conftest.py b/test_utils/managers/contest/conftest.py index a9375f6..b29cf18 100644 --- a/test_utils/managers/contest/conftest.py +++ b/test_utils/managers/contest/conftest.py @@ -1,10 +1,14 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest from generalresearch.managers.base import Permission from generalresearch.managers.thl.contest_manager import ContestManager -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.pg_helper import PostgresConfig @pytest.fixture(scope="session") diff --git a/test_utils/managers/gr/conftest.py b/test_utils/managers/gr/conftest.py index 40bd7b3..5392c69 100644 --- a/test_utils/managers/gr/conftest.py +++ b/test_utils/managers/gr/conftest.py @@ -3,6 +3,7 @@ from __future__ import annotations import subprocess from collections.abc import Callable, Generator from random import randint +from typing import TYPE_CHECKING import pytest import redis @@ -10,8 +11,6 @@ import redis.asyncio as redis_async from pydantic import PostgresDsn from redis import Redis -from generalresearch.config import GRLBaseSettings -from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager from generalresearch.managers.gr.business import ( BusinessAddressManager, BusinessBankAccountManager, @@ -20,6 +19,10 @@ from generalresearch.managers.gr.business import ( from generalresearch.pg_helper import PostgresConfig from generalresearch.redis_helper import RedisConfig +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager + # === Msc === @pytest.fixture(scope="session") diff --git a/test_utils/managers/ledger/conftest.py b/test_utils/managers/ledger/conftest.py index ce8348e..c60ee1b 100644 --- a/test_utils/managers/ledger/conftest.py +++ b/test_utils/managers/ledger/conftest.py @@ -1,18 +1,24 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest from generalresearch.managers.base import Permission from generalresearch.managers.thl.ledger_manager.ledger import ( - LedgerAccountManager, LedgerManager, - LedgerTransactionManager, ) from generalresearch.managers.thl.ledger_manager.thl_ledger import ( ThlLedgerManager, ) -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import ( + LedgerAccountManager, + LedgerTransactionManager, + ) + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig # --- Ledger --- diff --git a/test_utils/managers/thl/conftest.py b/test_utils/managers/thl/conftest.py index d40b7d2..af3fd23 100644 --- a/test_utils/managers/thl/conftest.py +++ b/test_utils/managers/thl/conftest.py @@ -1,44 +1,47 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING import pytest from pydantic import PostgresDsn -from generalresearch.config import GRLBaseSettings from generalresearch.managers.base import Permission -from generalresearch.managers.thl.buyer import BuyerManager -from generalresearch.managers.thl.category import CategoryManager -from generalresearch.managers.thl.payout import ( - BrokerageProductPayoutEventManager, - BusinessPayoutEventManager, - PayoutEventManager, - UserPayoutEventManager, -) -from generalresearch.managers.thl.product import ProductManager -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.task_adjustment import ( - TaskAdjustmentManager, -) from generalresearch.managers.thl.user_manager.mysql_user_manager import ( MysqlUserManager, ) from generalresearch.managers.thl.user_manager.redis_user_manager import ( RedisUserManager, ) -from generalresearch.managers.thl.user_manager.user_manager import ( - UserManager, -) -from generalresearch.managers.thl.user_manager.user_metadata_manager import ( - UserMetadataManager, -) -from generalresearch.managers.thl.wall import ( - WallCacheManager, - WallManager, -) from generalresearch.pg_helper import PostgresConfig from generalresearch.redis_helper import RedisConfig +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.managers.thl.buyer import BuyerManager + from generalresearch.managers.thl.category import CategoryManager + from generalresearch.managers.thl.payout import ( + BrokerageProductPayoutEventManager, + BusinessPayoutEventManager, + PayoutEventManager, + UserPayoutEventManager, + ) + from generalresearch.managers.thl.product import ProductManager + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.task_adjustment import ( + TaskAdjustmentManager, + ) + from generalresearch.managers.thl.user_manager.user_manager import ( + UserManager, + ) + from generalresearch.managers.thl.user_manager.user_metadata_manager import ( + UserMetadataManager, + ) + from generalresearch.managers.thl.wall import ( + WallCacheManager, + WallManager, + ) + @pytest.fixture(scope="session") def thl_web_rr(django_db_factory: Callable[..., PostgresDsn]) -> PostgresConfig: diff --git a/test_utils/managers/upk/conftest.py b/test_utils/managers/upk/conftest.py index 7eabee1..f581278 100644 --- a/test_utils/managers/upk/conftest.py +++ b/test_utils/managers/upk/conftest.py @@ -1,4 +1,5 @@ from collections.abc import Callable, Generator +from typing import TYPE_CHECKING import pytest @@ -12,9 +13,11 @@ from generalresearch.managers.thl.profiling.uqa import UQAManager from generalresearch.managers.thl.profiling.user_upk import ( UserUpkManager, ) -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig @pytest.fixture(scope="session") diff --git a/test_utils/models/conftest.py b/test_utils/models/conftest.py index 3a10ea3..089f2e6 100644 --- a/test_utils/models/conftest.py +++ b/test_utils/models/conftest.py @@ -12,13 +12,12 @@ import pytest from pydantic import AwareDatetime, PositiveInt from pytest import FixtureRequest as Request -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( WALL_ALLOWED_STATUS_STATUS_CODE, Status, ) from generalresearch.models.thl.survey.model import Buyer, Survey -from generalresearch.pg_helper import PostgresConfig if TYPE_CHECKING: from generalresearch.currency import USDCent @@ -53,6 +52,7 @@ if TYPE_CHECKING: from generalresearch.models.thl.user import User from generalresearch.models.thl.user_iphistory import IPRecord from generalresearch.models.thl.userhealth import AuditLog, AuditLogLevel + from generalresearch.pg_helper import PostgresConfig # === THL === diff --git a/test_utils/models/contest/conftest.py b/test_utils/models/contest/conftest.py index 84930b8..91425dc 100644 --- a/test_utils/models/contest/conftest.py +++ b/test_utils/models/contest/conftest.py @@ -3,36 +3,41 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pytest import FixtureRequest as Request from generalresearch.currency import USDCent -from generalresearch.managers.thl.contest_manager import ContestManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.contest import ( ContestEndCondition, ContestPrize, ) -from generalresearch.models.thl.contest.contest import Contest from generalresearch.models.thl.contest.definitions import ( ContestPrizeKind, ContestType, ) -from generalresearch.models.thl.contest.leaderboard import ( - LeaderboardContestCreate, -) -from generalresearch.models.thl.contest.milestone import ( - MilestoneContestCreate, -) from generalresearch.models.thl.contest.raffle import ( ContestEntryType, - RaffleContest, RaffleContestCreate, ) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.contest_manager import ContestManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.contest.contest import Contest + from generalresearch.models.thl.contest.leaderboard import ( + LeaderboardContestCreate, + ) + from generalresearch.models.thl.contest.milestone import ( + MilestoneContestCreate, + ) + from generalresearch.models.thl.contest.raffle import ( + RaffleContest, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User # === Miscellaneous === diff --git a/test_utils/models/gr/conftest.py b/test_utils/models/gr/conftest.py index b623255..6c1877a 100644 --- a/test_utils/models/gr/conftest.py +++ b/test_utils/models/gr/conftest.py @@ -2,30 +2,32 @@ from __future__ import annotations from collections.abc import Callable from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pydantic import PositiveInt from pydantic_extra_types.phone_numbers import PhoneNumber -from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager -from generalresearch.managers.gr.business import ( - BusinessAddressManager, - BusinessBankAccountManager, - BusinessManager, -) -from generalresearch.managers.gr.team import MembershipManager, TeamManager -from generalresearch.models.custom_types import UUIDStr -from generalresearch.models.gr.authentication import GRToken, GRUser -from generalresearch.models.gr.business import ( - Business, - BusinessAddress, - BusinessBankAccount, - TransferMethod, -) -from generalresearch.models.gr.team import Membership, Team -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig +if TYPE_CHECKING: + from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager + from generalresearch.managers.gr.business import ( + BusinessAddressManager, + BusinessBankAccountManager, + BusinessManager, + ) + from generalresearch.managers.gr.team import MembershipManager, TeamManager + from generalresearch.models.custom_types import UUIDStr + from generalresearch.models.gr.authentication import GRToken, GRUser + from generalresearch.models.gr.business import ( + Business, + BusinessAddress, + BusinessBankAccount, + TransferMethod, + ) + from generalresearch.models.gr.team import Membership, Team + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig # --- Static --- diff --git a/test_utils/models/ledger/conftest.py b/test_utils/models/ledger/conftest.py index 1c1027c..8437c7f 100644 --- a/test_utils/models/ledger/conftest.py +++ b/test_utils/models/ledger/conftest.py @@ -11,7 +11,6 @@ import pytest from pytest import FixtureRequest as Request from generalresearch.currency import USDCent -from generalresearch.managers.base import PostgresManager from test_utils.models.conftest import ( payout_config, product_amt_true, @@ -24,6 +23,9 @@ from test_utils.models.conftest import ( wall_factory, ) +if TYPE_CHECKING: + from generalresearch.managers.base import PostgresManager + _ = ( user_factory, product_user_wallet_no, diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py index 6ba37a3..4ff59ee 100644 --- a/test_utils/models/network/conftest.py +++ b/test_utils/models/network/conftest.py @@ -1,5 +1,6 @@ import os from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -23,7 +24,9 @@ from generalresearch.models.network.tool_run_command import ( RDNSRunCommand, RDNSRunCommandOptions, ) -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.pg_helper import PostgresConfig @pytest.fixture(scope="session") diff --git a/test_utils/models/thl/conftest.py b/test_utils/models/thl/conftest.py index fc57c73..3c77e27 100644 --- a/test_utils/models/thl/conftest.py +++ b/test_utils/models/thl/conftest.py @@ -5,7 +5,7 @@ from datetime import UTC, datetime from decimal import ROUND_DOWN, Decimal from random import choice as rand_choice from random import randint, random -from typing import Any +from typing import TYPE_CHECKING, Any from uuid import uuid4 import faker @@ -13,47 +13,53 @@ import pytest from grip_client.enums import AccessType from pydantic import PositiveInt -from generalresearch.managers.thl.ipinfo import IPGeonameManager, IPInformationManager -from generalresearch.managers.thl.payout import UserPayoutEventManager -from generalresearch.managers.thl.product import ProductManager -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.user_manager.user_manager import UserManager -from generalresearch.managers.thl.userhealth import AuditLogManager, IPRecordManager -from generalresearch.managers.thl.wall import WallManager -from generalresearch.models import DeviceType -from generalresearch.models.custom_types import ( - AwareDatetimeISO, - IPvAnyAddressStr, - UUIDStr, -) -from generalresearch.models.legacy.bucket import Bucket -from generalresearch.models.thl.definitions import ( - PayoutStatus, -) -from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation -from generalresearch.models.thl.payout import UserPayoutEvent -from generalresearch.models.thl.product import ( - PayoutConfig, - Product, - ProfilingConfig, - SessionConfig, - SourcesConfig, - SupplyConfig, - UserCreateConfig, - UserHealthConfig, - UserWalletConfig, -) +from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.session import ( - Session, Source, Status, - Wall, ) from generalresearch.models.thl.user import User -from generalresearch.models.thl.user_iphistory import IPRecord -from generalresearch.models.thl.userhealth import AuditLog, AuditLogLevel -from generalresearch.models.thl.wallet import PayoutType -from generalresearch.models.thl.wallet.cashout_method import CashMailOrderData +from generalresearch.models.thl.userhealth import AuditLogLevel +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.managers.thl.ipinfo import ( + IPGeonameManager, + IPInformationManager, + ) + from generalresearch.managers.thl.payout import UserPayoutEventManager + from generalresearch.managers.thl.product import ProductManager + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.user_manager.user_manager import UserManager + from generalresearch.managers.thl.userhealth import AuditLogManager, IPRecordManager + from generalresearch.managers.thl.wall import WallManager + from generalresearch.models.custom_types import ( + AwareDatetimeISO, + IPvAnyAddressStr, + UUIDStr, + ) + from generalresearch.models.definitions import DeviceType + from generalresearch.models.legacy.bucket import Bucket + from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation + from generalresearch.models.thl.payout import UserPayoutEvent + from generalresearch.models.thl.product import ( + PayoutConfig, + Product, + ProfilingConfig, + SessionConfig, + SourcesConfig, + SupplyConfig, + UserCreateConfig, + UserHealthConfig, + UserWalletConfig, + ) + from generalresearch.models.thl.session import ( + Session, + Wall, + ) + from generalresearch.models.thl.user_iphistory import IPRecord + from generalresearch.models.thl.userhealth import AuditLog + from generalresearch.models.thl.wallet.cashout_method import CashMailOrderData fake = faker.Faker() diff --git a/test_utils/models/upk/conftest.py b/test_utils/models/upk/conftest.py index ef77dd6..59266b2 100644 --- a/test_utils/models/upk/conftest.py +++ b/test_utils/models/upk/conftest.py @@ -9,10 +9,9 @@ from uuid import UUID import pandas as pd import pytest -from generalresearch.pg_helper import PostgresConfig - if TYPE_CHECKING: from generalresearch.managers.thl.category import CategoryManager + from generalresearch.pg_helper import PostgresConfig def insert_data_from_csv( diff --git a/test_utils/spectrum/conftest.py b/test_utils/spectrum/conftest.py index a8ce9d9..cc91cff 100644 --- a/test_utils/spectrum/conftest.py +++ b/test_utils/spectrum/conftest.py @@ -3,16 +3,15 @@ from __future__ import annotations import time from datetime import UTC, datetime from decimal import Decimal -from typing import Any +from typing import TYPE_CHECKING, Any import pytest -from generalresearch.config import GRLBaseSettings from generalresearch.managers.spectrum.survey import ( SpectrumCriteriaManager, SpectrumSurveyManager, ) -from generalresearch.models import ( +from generalresearch.models.definitions import ( LogicalOperator, ) from generalresearch.models.spectrum.survey import ( @@ -22,6 +21,9 @@ from generalresearch.models.spectrum.survey import ( from generalresearch.models.thl.survey.condition import ConditionValueType from generalresearch.sql_helper import SqlHelper +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + @pytest.fixture(scope="session") def spectrum_rw(settings: GRLBaseSettings) -> SqlHelper: diff --git a/tests/incite/collections/test_df_collection_base.py b/tests/incite/collections/test_df_collection_base.py index c236700..e20b44b 100644 --- a/tests/incite/collections/test_df_collection_base.py +++ b/tests/incite/collections/test_df_collection_base.py @@ -9,10 +9,10 @@ from generalresearch.incite.collections import ( DFCollection, DFCollectionType, ) -from generalresearch.pg_helper import PostgresConfig if TYPE_CHECKING: from generalresearch.incite.base import GRLDatasets + from generalresearch.pg_helper import PostgresConfig df_collection_types = [e for e in DFCollectionType if e is not DFCollectionType.TEST] diff --git a/tests/incite/collections/test_df_collection_item_base.py b/tests/incite/collections/test_df_collection_item_base.py index e0171c2..fd70bf0 100644 --- a/tests/incite/collections/test_df_collection_item_base.py +++ b/tests/incite/collections/test_df_collection_item_base.py @@ -5,15 +5,16 @@ from typing import TYPE_CHECKING import pytest -from generalresearch.incite.collections import ( +from generalresearch.incite.collections.base import ( DFCollection, DFCollectionItem, DFCollectionType, ) -from generalresearch.pg_helper import PostgresConfig if TYPE_CHECKING: from generalresearch.incite.base import GRLDatasets + from generalresearch.pg_helper import PostgresConfig + df_collection_types = [e for e in DFCollectionType if e is not DFCollectionType.TEST] diff --git a/tests/incite/collections/test_df_collection_item_thl_web.py b/tests/incite/collections/test_df_collection_item_thl_web.py index 5f9a3f6..061c576 100644 --- a/tests/incite/collections/test_df_collection_item_thl_web.py +++ b/tests/incite/collections/test_df_collection_item_thl_web.py @@ -5,6 +5,7 @@ from datetime import UTC, datetime, timedelta from itertools import product as iter_product from os.path import join as pjoin from pathlib import Path, PurePath +from typing import TYPE_CHECKING from uuid import uuid4 import dask.dataframe as dd @@ -21,19 +22,24 @@ from faker import Faker from pandera.pandas import DataFrameSchema from pydantic import FilePath -from generalresearch.incite.base import CollectionItemBase, GRLDatasets -from generalresearch.incite.collections import ( - DFCollection, - DFCollectionItem, +from generalresearch.incite.base import CollectionItemBase +from generalresearch.incite.collections.base import ( DFCollectionType, ) from generalresearch.incite.schemas import ARCHIVE_AFTER -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User from generalresearch.pg_helper import PostgresConfig from generalresearch.sql_helper import PostgresDsn +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.base import ( + DFCollection, + DFCollectionItem, + ) + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + fake = Faker() df_collections = [ diff --git a/tests/incite/collections/test_df_collection_thl_marketplaces.py b/tests/incite/collections/test_df_collection_thl_marketplaces.py index 6a7e5c9..6ad0cb4 100644 --- a/tests/incite/collections/test_df_collection_thl_marketplaces.py +++ b/tests/incite/collections/test_df_collection_thl_marketplaces.py @@ -1,18 +1,21 @@ from datetime import UTC, datetime from itertools import product +from typing import TYPE_CHECKING import pytest from pandera.pandas import Column, DataFrameSchema, Index -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections import DFCollection, DFCollectionType +from generalresearch.incite.collections.base import DFCollection, DFCollectionType from generalresearch.incite.collections.thl_marketplaces import ( InnovateSurveyHistoryCollection, MorningSurveyTimeseriesCollection, SagoSurveyHistoryCollection, SpectrumSurveyTimeseriesCollection, ) -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.pg_helper import PostgresConfig def combo_object(): diff --git a/tests/incite/collections/test_df_collection_thl_web.py b/tests/incite/collections/test_df_collection_thl_web.py index 6d509bc..20d7187 100644 --- a/tests/incite/collections/test_df_collection_thl_web.py +++ b/tests/incite/collections/test_df_collection_thl_web.py @@ -9,7 +9,7 @@ import pandas as pd import pytest from pandera.pandas import DataFrameSchema -from generalresearch.incite.collections import ( +from generalresearch.incite.collections.base import ( DFCollection, DFCollectionType, ) diff --git a/tests/incite/mergers/foundations/test_enriched_session.py b/tests/incite/mergers/foundations/test_enriched_session.py index 2a161e4..71b2442 100644 --- a/tests/incite/mergers/foundations/test_enriched_session.py +++ b/tests/incite/mergers/foundations/test_enriched_session.py @@ -4,29 +4,32 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal from itertools import product +from typing import TYPE_CHECKING import dask.dataframe as dd import pandas as pd import pytest from dask.distributed import Client as DaskClient -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - WallDFCollection, -) -from generalresearch.incite.mergers.foundations.enriched_session import ( - EnrichedSessionMerge, -) from generalresearch.incite.schemas.admin_responses import ( AdminPOPSessionSchema, ) -from generalresearch.models.admin.request import ( - ReportRequest, -) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + WallDFCollection, + ) + from generalresearch.incite.mergers.foundations.enriched_session import ( + EnrichedSessionMerge, + ) + from generalresearch.models.admin.request import ( + ReportRequest, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig @pytest.mark.parametrize( diff --git a/tests/incite/mergers/foundations/test_enriched_task_adjust.py b/tests/incite/mergers/foundations/test_enriched_task_adjust.py index 0606b6f..877d22f 100644 --- a/tests/incite/mergers/foundations/test_enriched_task_adjust.py +++ b/tests/incite/mergers/foundations/test_enriched_task_adjust.py @@ -3,26 +3,28 @@ from __future__ import annotations from collections.abc import Callable from datetime import timedelta from itertools import product as iter_product +from typing import TYPE_CHECKING import dask.dataframe as dd import pandas as pd import pytest from dask.distributed import Client as DaskClient -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - TaskAdjustmentDFCollection, - WallDFCollection, -) -from generalresearch.incite.mergers.foundations.enriched_task_adjust import ( - EnrichedTaskAdjustMerge, -) -from generalresearch.incite.mergers.foundations.enriched_wall import ( - EnrichedWallMerge, -) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig +if TYPE_CHECKING: + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + TaskAdjustmentDFCollection, + WallDFCollection, + ) + from generalresearch.incite.mergers.foundations.enriched_task_adjust import ( + EnrichedTaskAdjustMerge, + ) + from generalresearch.incite.mergers.foundations.enriched_wall import ( + EnrichedWallMerge, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig @pytest.mark.parametrize( diff --git a/tests/incite/mergers/foundations/test_enriched_wall.py b/tests/incite/mergers/foundations/test_enriched_wall.py index 0cb8f60..2b9afb8 100644 --- a/tests/incite/mergers/foundations/test_enriched_wall.py +++ b/tests/incite/mergers/foundations/test_enriched_wall.py @@ -2,27 +2,32 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal from itertools import product as iter_product +from typing import TYPE_CHECKING import dask.dataframe as dd import pandas as pd import pytest from dask.distributed import Client as DaskClient -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - WallDFCollection, -) - -# noinspection PyUnresolvedReferences from generalresearch.incite.mergers.foundations.enriched_wall import ( - EnrichedWallMerge, EnrichedWallMergeItem, ) -from generalresearch.models.admin.request import ReportRequest -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + WallDFCollection, + ) + + # noinspection PyUnresolvedReferences + from generalresearch.incite.mergers.foundations.enriched_wall import ( + EnrichedWallMerge, + ) + from generalresearch.models.admin.request import ReportRequest + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig @pytest.mark.parametrize( diff --git a/tests/incite/mergers/foundations/test_user_id_product.py b/tests/incite/mergers/foundations/test_user_id_product.py index 7367056..8c4b2f7 100644 --- a/tests/incite/mergers/foundations/test_user_id_product.py +++ b/tests/incite/mergers/foundations/test_user_id_product.py @@ -2,17 +2,22 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta from itertools import product +from typing import TYPE_CHECKING import pandas as pd import pytest from dask.distributed import Client as DaskClient -# noinspection PyUnresolvedReferences from generalresearch.incite.mergers.foundations.user_id_product import ( - UserIdProductMerge, UserIdProductMergeItem, ) +if TYPE_CHECKING: + # noinspection PyUnresolvedReferences + from generalresearch.incite.mergers.foundations.user_id_product import ( + UserIdProductMerge, + ) + @pytest.mark.parametrize( argnames="offset, duration, start", diff --git a/tests/incite/mergers/test_merge_collection.py b/tests/incite/mergers/test_merge_collection.py index cf8315f..3f478bd 100644 --- a/tests/incite/mergers/test_merge_collection.py +++ b/tests/incite/mergers/test_merge_collection.py @@ -2,17 +2,20 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta from itertools import product +from typing import TYPE_CHECKING import pandas as pd import pytest from pandera.pandas import DataFrameSchema -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.mergers import ( +from generalresearch.incite.mergers.base import ( MergeCollection, MergeType, ) +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + merge_types = [e for e in MergeType if e != MergeType.TEST] diff --git a/tests/incite/mergers/test_merge_collection_item.py b/tests/incite/mergers/test_merge_collection_item.py index 5ca2f6b..baf1bc4 100644 --- a/tests/incite/mergers/test_merge_collection_item.py +++ b/tests/incite/mergers/test_merge_collection_item.py @@ -3,14 +3,17 @@ from __future__ import annotations from datetime import timedelta from itertools import product from pathlib import PurePath +from typing import TYPE_CHECKING import pytest -from generalresearch.incite.mergers import ( - MergeCollection, - MergeCollectionItem, - MergeType, -) +from generalresearch.incite.mergers.base import MergeType + +if TYPE_CHECKING: + from generalresearch.incite.mergers.base import ( + MergeCollection, + MergeCollectionItem, + ) @pytest.mark.parametrize( diff --git a/tests/incite/mergers/test_pop_ledger.py b/tests/incite/mergers/test_pop_ledger.py index 2146344..9ec188b 100644 --- a/tests/incite/mergers/test_pop_ledger.py +++ b/tests/incite/mergers/test_pop_ledger.py @@ -3,23 +3,26 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta from itertools import product as iter_product +from typing import TYPE_CHECKING import pandas as pd import pytest from dask.distributed import Client as DaskClient -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections.thl_web import ( - LedgerDFCollection, - SessionDFCollection, -) -from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.thl_web import ( + LedgerDFCollection, + SessionDFCollection, + ) + from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User @pytest.mark.parametrize( diff --git a/tests/incite/mergers/test_ym_survey_merge.py b/tests/incite/mergers/test_ym_survey_merge.py index 8a4897b..d83a98c 100644 --- a/tests/incite/mergers/test_ym_survey_merge.py +++ b/tests/incite/mergers/test_ym_survey_merge.py @@ -3,22 +3,24 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta from itertools import product +from typing import TYPE_CHECKING import pandas as pd import pytest from dask.distributed import Client as DaskClient -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - WallDFCollection, -) -from generalresearch.incite.mergers.foundations.enriched_session import ( - EnrichedSessionMerge, -) -from generalresearch.incite.mergers.ym_survey_wall import YMSurveyWallMerge -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig +if TYPE_CHECKING: + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + WallDFCollection, + ) + from generalresearch.incite.mergers.foundations.enriched_session import ( + EnrichedSessionMerge, + ) + from generalresearch.incite.mergers.ym_survey_wall import YMSurveyWallMerge + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig # noinspection PyUnresolvedReferences diff --git a/tests/incite/test_collection_base.py b/tests/incite/test_collection_base.py index 577eda9..1a664a2 100644 --- a/tests/incite/test_collection_base.py +++ b/tests/incite/test_collection_base.py @@ -4,6 +4,7 @@ from datetime import UTC, datetime, timedelta, timezone from os.path import exists as pexists from os.path import join as pjoin from pathlib import Path +from typing import TYPE_CHECKING from uuid import uuid4 import numpy as np @@ -11,7 +12,10 @@ import pandas as pd import pytest from _pytest._code.code import ExceptionInfo -from generalresearch.incite.base import CollectionBase, GRLDatasets +from generalresearch.incite.base import CollectionBase + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets AGO_15min = (datetime.now(tz=UTC) - timedelta(minutes=15)).replace(microsecond=0) AGO_1HR = (datetime.now(tz=UTC) - timedelta(hours=1)).replace(microsecond=0) diff --git a/tests/incite/test_collection_base_item.py b/tests/incite/test_collection_base_item.py index e09f54a..b9f1c26 100644 --- a/tests/incite/test_collection_base_item.py +++ b/tests/incite/test_collection_base_item.py @@ -3,6 +3,7 @@ from __future__ import annotations from datetime import UTC, datetime from os.path import join as pjoin from pathlib import Path +from typing import TYPE_CHECKING from uuid import uuid4 import dask.dataframe as dd @@ -10,7 +11,10 @@ import pandas as pd import pytest from pydantic import ValidationError -from generalresearch.incite.base import CollectionItemBase, GRLDatasets +from generalresearch.incite.base import CollectionItemBase + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets class TestCollectionItemBase: diff --git a/tests/managers/gr/test_business.py b/tests/managers/gr/test_business.py index ed141b1..1a5d4fa 100644 --- a/tests/managers/gr/test_business.py +++ b/tests/managers/gr/test_business.py @@ -1,21 +1,24 @@ +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.gr.business import ( - BusinessAddressManager, - BusinessBankAccountManager, - BusinessManager, -) -from generalresearch.managers.gr.team import MembershipManager, TeamManager -from generalresearch.models.gr.authentication import GRUser from generalresearch.models.gr.business import ( Business, BusinessAddress, BusinessBankAccount, TransferMethod, ) -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.managers.gr.business import ( + BusinessAddressManager, + BusinessBankAccountManager, + BusinessManager, + ) + from generalresearch.managers.gr.team import MembershipManager, TeamManager + from generalresearch.models.gr.authentication import GRUser + from generalresearch.pg_helper import PostgresConfig class TestBusinessBankAccountManager: diff --git a/tests/managers/gr/test_team.py b/tests/managers/gr/test_team.py index ae3e1bb..17e0470 100644 --- a/tests/managers/gr/test_team.py +++ b/tests/managers/gr/test_team.py @@ -1,15 +1,18 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 -from generalresearch.managers.gr.authentication import GRUserManager -from generalresearch.managers.gr.team import MembershipManager, TeamManager -from generalresearch.models.gr.authentication import GRUser from generalresearch.models.gr.team import Membership, Team -from generalresearch.models.thl.product import Product -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.managers.gr.authentication import GRUserManager + from generalresearch.managers.gr.team import MembershipManager, TeamManager + from generalresearch.models.gr.authentication import GRUser + from generalresearch.models.thl.product import Product + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig class TestMembershipManager: diff --git a/tests/managers/leaderboard.py b/tests/managers/leaderboard.py index d97714d..197477b 100644 --- a/tests/managers/leaderboard.py +++ b/tests/managers/leaderboard.py @@ -6,6 +6,7 @@ import zoneinfo from collections.abc import Callable from datetime import UTC, datetime from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -26,7 +27,9 @@ from generalresearch.models.thl.product import ( ) from generalresearch.models.thl.session import Session from generalresearch.models.thl.user import User -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.redis_helper import RedisConfig # random uuid for leaderboard tests product_id = uuid4().hex diff --git a/tests/managers/network/test_label.py b/tests/managers/network/test_label.py index 71efa95..abdd28f 100644 --- a/tests/managers/network/test_label.py +++ b/tests/managers/network/test_label.py @@ -1,11 +1,12 @@ import ipaddress +from datetime import datetime +from typing import TYPE_CHECKING import faker import pytest from psycopg.errors import UniqueViolation from pydantic import ValidationError -from generalresearch.managers.network.label import IPLabelManager from generalresearch.models.network.label import ( IPLabel, IPLabelKind, @@ -14,11 +15,14 @@ from generalresearch.models.network.label import ( ) from generalresearch.models.thl.ipinfo import normalize_ip +if TYPE_CHECKING: + from generalresearch.managers.network.label import IPLabelManager + fake = faker.Faker() @pytest.fixture -def ip_label(utc_now) -> IPLabel: +def ip_label(utc_now: datetime) -> IPLabel: ip = ipaddress.IPv6Network((fake.ipv6(), 64), strict=False) return IPLabel( label_kind=IPLabelKind.VPN, @@ -31,7 +35,7 @@ def ip_label(utc_now) -> IPLabel: ) -def test_model(utc_now): +def test_model(utc_now: datetime): ip = fake.ipv4_public() lbl = IPLabel( label_kind=IPLabelKind.VPN, @@ -142,7 +146,7 @@ def test_filter_network( assert len(res) == 2 -def test_network(iplabel_manager: IPLabelManager, utc_now): +def test_network(iplabel_manager: IPLabelManager, utc_now: datetime): # This is a fully-specific /128 ipv6 address. # e.g. '51b7:b38d:8717:6c5b:cd3e:f5c3:3aba:17d' ip = fake.ipv6() @@ -174,7 +178,10 @@ def test_network(iplabel_manager: IPLabelManager, utc_now): def test_label_cidr_and_ipinfo( - iplabel_manager: IPLabelManager, ip_information_factory, ip_geoname, utc_now + iplabel_manager: IPLabelManager, + ip_information_factory, + ip_geoname, + utc_now: datetime, ): # We have network_iplabel.ip as a cidr col and # thl_ipinformation.ip as a inet col. Make sure we can join appropriately diff --git a/tests/managers/test_events.py b/tests/managers/test_events.py index cb32275..8745126 100644 --- a/tests/managers/test_events.py +++ b/tests/managers/test_events.py @@ -8,13 +8,13 @@ from datetime import UTC, datetime, timedelta from decimal import Decimal from functools import partial from math import floor +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.events import EventManager, EventSubscriber -from generalresearch.managers.thl.product import ProductManager -from generalresearch.models import Source +from generalresearch.managers.events import EventSubscriber +from generalresearch.models.definitions import Source from generalresearch.models.events import ( AggregateBySource, EventType, @@ -25,7 +25,11 @@ from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.definitions import Status, StatusCode1 from generalresearch.models.thl.session import Session, Wall from generalresearch.models.thl.user import User -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.managers.events import EventManager + from generalresearch.managers.thl.product import ProductManager + from generalresearch.redis_helper import RedisConfig # We don't need anything in the db, so not using the db fixtures diff --git a/tests/managers/test_lucid.py b/tests/managers/test_lucid.py index 20dca22..6771a0c 100644 --- a/tests/managers/test_lucid.py +++ b/tests/managers/test_lucid.py @@ -1,9 +1,13 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest from generalresearch.managers.lucid.profiling import get_profiling_library -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.pg_helper import PostgresConfig qids = ["42", "43", "45", "97", "120", "639", "15297"] diff --git a/tests/managers/thl/test_buyer.py b/tests/managers/thl/test_buyer.py index 6776ab3..0ab2d52 100644 --- a/tests/managers/thl/test_buyer.py +++ b/tests/managers/thl/test_buyer.py @@ -1,9 +1,12 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING -from generalresearch.managers.thl.buyer import BuyerManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source + +if TYPE_CHECKING: + from generalresearch.managers.thl.buyer import BuyerManager class TestBuyer: diff --git a/tests/managers/thl/test_cashout_method.py b/tests/managers/thl/test_cashout_method.py index ca85c6b..877d7b2 100644 --- a/tests/managers/thl/test_cashout_method.py +++ b/tests/managers/thl/test_cashout_method.py @@ -1,21 +1,26 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING import pytest -from generalresearch.config import GRLBaseSettings -from generalresearch.managers.thl.cashout_method import ( - CashoutMethodManager, -) -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType from generalresearch.models.thl.wallet.cashout_method import ( CashMailCashoutMethodData, - CashoutMethod, PaypalCashoutMethodData, USDeliveryAddress, ) +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.managers.thl.cashout_method import ( + CashoutMethodManager, + ) + from generalresearch.models.thl.user import User + from generalresearch.models.thl.wallet.cashout_method import ( + CashoutMethod, + ) class TestTangoCashoutMethods: diff --git a/tests/managers/thl/test_category.py b/tests/managers/thl/test_category.py index ec52aae..4d00643 100644 --- a/tests/managers/thl/test_category.py +++ b/tests/managers/thl/test_category.py @@ -1,12 +1,15 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.category import CategoryManager from generalresearch.models.thl.category import Category -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.managers.thl.category import CategoryManager + from generalresearch.pg_helper import PostgresConfig class TestCategory: diff --git a/tests/managers/thl/test_contest/test_leaderboard.py b/tests/managers/thl/test_contest/test_leaderboard.py index 3a63075..d80d512 100644 --- a/tests/managers/thl/test_contest/test_leaderboard.py +++ b/tests/managers/thl/test_contest/test_leaderboard.py @@ -1,23 +1,28 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING from zoneinfo import ZoneInfo from generalresearch.currency import USDCent -from generalresearch.managers.thl.contest_manager import ContestManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.user_manager.user_manager import UserManager from generalresearch.models.thl.contest.definitions import ( ContestEndReason, ContestStatus, ) from generalresearch.models.thl.contest.leaderboard import ( LeaderboardContest, - LeaderboardContestCreate, ) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.managers.thl.contest_manager import ContestManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.user_manager.user_manager import UserManager + from generalresearch.models.thl.contest.leaderboard import ( + LeaderboardContestCreate, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + from generalresearch.redis_helper import RedisConfig class TestLeaderboardContestCRUD: diff --git a/tests/managers/thl/test_contest/test_milestone.py b/tests/managers/thl/test_contest/test_milestone.py index e29ba4c..dbb2016 100644 --- a/tests/managers/thl/test_contest/test_milestone.py +++ b/tests/managers/thl/test_contest/test_milestone.py @@ -2,9 +2,8 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime +from typing import TYPE_CHECKING -from generalresearch.managers.thl.contest_manager import ContestManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.contest.definitions import ( ContestEndReason, ContestStatus, @@ -12,12 +11,18 @@ from generalresearch.models.thl.contest.definitions import ( from generalresearch.models.thl.contest.milestone import ( ContestEntryTrigger, MilestoneContest, - MilestoneContestCreate, MilestoneUserView, ) -from generalresearch.models.thl.contest.raffle import RaffleContest -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.contest_manager import ContestManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.contest.milestone import ( + MilestoneContestCreate, + ) + from generalresearch.models.thl.contest.raffle import RaffleContest + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User class TestMilestoneContest: @@ -25,8 +30,6 @@ class TestMilestoneContest: def test_should_end( self, contest: MilestoneContest, - thl_ledger_manager: ThlLedgerManager, - contest_manager: ContestManager, ): # contest is active and has no entries should, msg = contest.should_end() @@ -53,7 +56,6 @@ class TestMilestoneContestCRUD: self, contest_create: MilestoneContestCreate, product_user_wallet_yes: Product, - thl_ledger_manager: ThlLedgerManager, contest_manager: ContestManager, ): c = contest_manager.create( diff --git a/tests/managers/thl/test_contest/test_raffle.py b/tests/managers/thl/test_contest/test_raffle.py index 06d4676..7803952 100644 --- a/tests/managers/thl/test_contest/test_raffle.py +++ b/tests/managers/thl/test_contest/test_raffle.py @@ -2,19 +2,17 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime +from typing import TYPE_CHECKING import pytest from pydantic import ValidationError from pytest import approx from generalresearch.currency import USDCent -from generalresearch.managers.thl.contest_manager import ContestManager from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionConditionFailedError, ) -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.contest import ( - Contest, ContestEndCondition, ContestEntryRule, ContestPrize, @@ -29,11 +27,20 @@ from generalresearch.models.thl.contest.raffle import ( ContestEntry, ContestEntryType, RaffleContest, - RaffleContestCreate, - RaffleUserView, ) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.contest_manager import ContestManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.contest import ( + Contest, + ) + from generalresearch.models.thl.contest.raffle import ( + RaffleContestCreate, + RaffleUserView, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User class TestRaffleContest: diff --git a/tests/managers/thl/test_harmonized_uqa.py b/tests/managers/thl/test_harmonized_uqa.py index 84eeb56..2fc0ff0 100644 --- a/tests/managers/thl/test_harmonized_uqa.py +++ b/tests/managers/thl/test_harmonized_uqa.py @@ -1,15 +1,18 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.profiling.uqa import UQAManager from generalresearch.models.thl.profiling.user_question_answer import ( DUMMY_UQA, UserQuestionAnswer, ) -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.profiling.uqa import UQAManager + from generalresearch.models.thl.user import User @pytest.mark.usefixtures("uqa_db_index", "upk_data", "uqa_manager_clear_cache") diff --git a/tests/managers/thl/test_ipinfo.py b/tests/managers/thl/test_ipinfo.py index 48b9efd..6954163 100644 --- a/tests/managers/thl/test_ipinfo.py +++ b/tests/managers/thl/test_ipinfo.py @@ -1,4 +1,5 @@ from collections.abc import Callable +from typing import TYPE_CHECKING import faker @@ -12,8 +13,10 @@ from generalresearch.models.thl.ipinfo import ( IPGeoname, IPInformation, ) -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig fake = faker.Faker() diff --git a/tests/managers/thl/test_ledger/test_lm_accounts.py b/tests/managers/thl/test_ledger/test_lm_accounts.py index 7b65b2d..f5ed883 100644 --- a/tests/managers/thl/test_ledger/test_lm_accounts.py +++ b/tests/managers/thl/test_ledger/test_lm_accounts.py @@ -2,6 +2,7 @@ from __future__ import annotations from itertools import product as iproduct from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -13,13 +14,18 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerAccountDoesntExistError, ) from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.models.custom_types import AccountType, Direction, UUIDStr +from generalresearch.models.custom_types import AccountType, Direction from generalresearch.models.thl.ledger import ( LedgerAccount, LedgerEntry, - LedgerTransaction, ) +if TYPE_CHECKING: + from generalresearch.models.custom_types import UUIDStr + from generalresearch.models.thl.ledger import ( + LedgerTransaction, + ) + @pytest.mark.parametrize( argnames="currency, kind, acct_id", diff --git a/tests/managers/thl/test_ledger/test_lm_tx.py b/tests/managers/thl/test_ledger/test_lm_tx.py index ce609d6..445405e 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx.py +++ b/tests/managers/thl/test_ledger/test_lm_tx.py @@ -2,6 +2,7 @@ from __future__ import annotations from decimal import Decimal from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -12,11 +13,15 @@ from generalresearch.managers.thl.ledger_manager.ledger import ( ) from generalresearch.models.thl.ledger import ( Direction, - LedgerAccount, LedgerEntry, LedgerTransaction, ) +if TYPE_CHECKING: + from generalresearch.models.thl.ledger import ( + LedgerAccount, + ) + class TestLedgerManagerCreateTx: diff --git a/tests/managers/thl/test_ledger/test_lm_tx_entries.py b/tests/managers/thl/test_ledger/test_lm_tx_entries.py index 9925b87..03c6e02 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx_entries.py +++ b/tests/managers/thl/test_ledger/test_lm_tx_entries.py @@ -1,11 +1,17 @@ from __future__ import annotations -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager +from typing import TYPE_CHECKING + from generalresearch.models.thl.ledger import ( LedgerEntry, - LedgerTransaction, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.models.thl.ledger import ( + LedgerTransaction, + ) + class TestLedgerEntryManager: diff --git a/tests/managers/thl/test_ledger/test_lm_tx_locks.py b/tests/managers/thl/test_ledger/test_lm_tx_locks.py index e603632..9ecc1bc 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx_locks.py +++ b/tests/managers/thl/test_ledger/test_lm_tx_locks.py @@ -5,10 +5,10 @@ from collections.abc import Callable, Generator from datetime import UTC, datetime, timedelta from decimal import Decimal from logging import LogCaptureFixture +from typing import TYPE_CHECKING import pytest -from generalresearch.currency import LedgerCurrency from generalresearch.managers.thl.ledger_manager.conditions import ( generate_condition_mp_payment, ) @@ -17,11 +17,8 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionCreateLockError, LedgerTransactionFlagAlreadyExistsError, ) -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.ledger import LedgerTransaction -from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import ( Session, Status, @@ -29,7 +26,13 @@ from generalresearch.models.thl.session import ( Wall, WallAdjustedStatus, ) -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.currency import LedgerCurrency + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User logger = logging.getLogger("LedgerManager") diff --git a/tests/managers/thl/test_ledger/test_lm_tx_metadata.py b/tests/managers/thl/test_ledger/test_lm_tx_metadata.py index f63efa4..3d8cf89 100644 --- a/tests/managers/thl/test_ledger/test_lm_tx_metadata.py +++ b/tests/managers/thl/test_ledger/test_lm_tx_metadata.py @@ -1,9 +1,12 @@ from __future__ import annotations -from generalresearch.managers.thl.ledger_manager.ledger import ( - LedgerManager, - LedgerTransaction, -) +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import ( + LedgerManager, + LedgerTransaction, + ) class TestLedgerMetadataManager: diff --git a/tests/managers/thl/test_ledger/test_thl_lm_accounts.py b/tests/managers/thl/test_ledger/test_thl_lm_accounts.py index 60eb71c..adff446 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_accounts.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_accounts.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -9,18 +10,20 @@ from generalresearch.currency import LedgerCurrency from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerAccountDoesntExistError, ) -from generalresearch.managers.thl.ledger_manager.ledger import ( - LedgerAccountManager, - LedgerManager, -) -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.ledger import ( AccountType, Direction, LedgerAccount, ) from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import ( + LedgerAccountManager, + LedgerManager, + ) + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.user import User class TestThlLedgerManagerAccounts: diff --git a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py index b518453..14c5270 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_bp_payout.py @@ -5,6 +5,7 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -12,7 +13,7 @@ import redis from pydantic import RedisDsn from redis.lock import Lock -from generalresearch.currency import LedgerCurrency, USDCent +from generalresearch.currency import USDCent from generalresearch.managers.base import Permission from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionConditionFailedError, @@ -22,24 +23,27 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( ) from generalresearch.managers.thl.ledger_manager.ledger import LedgerTransaction from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.payout import ( - BrokerageProductPayoutEventManager, -) -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.ledger import Direction, TransactionType -from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import ( Session, Status, StatusCode1, Wall, ) -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType -from generalresearch.pg_helper import PostgresConfig +from generalresearch.models.thl.wallet.definitions import PayoutType from generalresearch.redis_helper import RedisConfig +if TYPE_CHECKING: + from generalresearch.currency import LedgerCurrency + from generalresearch.managers.thl.payout import ( + BrokerageProductPayoutEventManager, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + def broken_acquire(self, *args, **kwargs): raise redis.exceptions.TimeoutError("Simulated timeout during acquire") diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx.py b/tests/managers/thl/test_ledger/test_thl_lm_tx.py index 1860d6d..2e4ab5e 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx.py @@ -5,26 +5,21 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.currency import LedgerCurrency, USDCent +from generalresearch.currency import USDCent from generalresearch.managers.thl.ledger_manager.ledger import ( - LedgerManager, LedgerTransaction, ) -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.product import ProductManager -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.wall import WallManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( WALL_ALLOWED_STATUS_STATUS_CODE, ) from generalresearch.models.thl.ledger import ( Direction, - LedgerAccount, TransactionType, ) from generalresearch.models.thl.payout import UserPayoutEvent @@ -41,8 +36,21 @@ from generalresearch.models.thl.session import ( Wall, WallAdjustedStatus, ) -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.currency import LedgerCurrency + from generalresearch.managers.thl.ledger_manager.ledger import ( + LedgerManager, + ) + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.product import ProductManager + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.wall import WallManager + from generalresearch.models.thl.ledger import ( + LedgerAccount, + ) + from generalresearch.models.thl.user import User logger = logging.getLogger("LedgerManager") diff --git a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py index 82dc143..3fd21dc 100644 --- a/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py +++ b/tests/managers/thl/test_ledger/test_thl_lm_tx__user_payouts.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging from collections.abc import Callable from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -11,12 +12,14 @@ from generalresearch.managers.thl.ledger_manager.exceptions import ( LedgerTransactionConditionFailedError, LedgerTransactionFlagAlreadyExistsError, ) -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.payout import UserPayoutEvent -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User class TestLedgerManagerAMT: diff --git a/tests/managers/thl/test_ledger/test_thl_pem.py b/tests/managers/thl/test_ledger/test_thl_pem.py index 29341cf..9dbec48 100644 --- a/tests/managers/thl/test_ledger/test_thl_pem.py +++ b/tests/managers/thl/test_ledger/test_thl_pem.py @@ -2,28 +2,31 @@ from __future__ import annotations from collections.abc import Callable from random import randint +from typing import TYPE_CHECKING from uuid import UUID, uuid4 import pytest from generalresearch.currency import USDCent -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ( - ThlLedgerManager, -) -from generalresearch.managers.thl.payout import ( - BrokerageProductPayoutEventManager, - UserPayoutEventManager, -) from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.payout import ( BrokerageProductPayoutEvent, ) -from generalresearch.models.thl.product import Product from generalresearch.models.thl.wallet.cashout_method import ( CashoutRequestInfo, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ( + ThlLedgerManager, + ) + from generalresearch.managers.thl.payout import ( + BrokerageProductPayoutEventManager, + UserPayoutEventManager, + ) + from generalresearch.models.thl.product import Product + class TestThlPayoutEventManager: diff --git a/tests/managers/thl/test_ledger/test_user_txs.py b/tests/managers/thl/test_ledger/test_user_txs.py index 56dc485..1c08498 100644 --- a/tests/managers/thl/test_ledger/test_user_txs.py +++ b/tests/managers/thl/test_ledger/test_user_txs.py @@ -3,12 +3,9 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 -from generalresearch.config import GRLBaseSettings -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.payout import UserPayoutEventManager from generalresearch.managers.thl.user_compensate import user_compensate from generalresearch.models.thl.definitions import ( Status, @@ -18,10 +15,16 @@ from generalresearch.models.thl.ledger import ( UserLedgerTransactionTypesSummary, UserLedgerTransactionTypeSummary, ) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.payout import UserPayoutEventManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User def test_user_txs( diff --git a/tests/managers/thl/test_ledger/test_wallet.py b/tests/managers/thl/test_ledger/test_wallet.py index cad3ea4..1ee9bf9 100644 --- a/tests/managers/thl/test_ledger/test_wallet.py +++ b/tests/managers/thl/test_ledger/test_wallet.py @@ -2,12 +2,11 @@ from __future__ import annotations from collections.abc import Callable from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.product import ProductManager from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, @@ -15,7 +14,11 @@ from generalresearch.models.thl.product import ( Product, UserWalletConfig, ) -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.product import ProductManager + from generalresearch.models.thl.user import User @pytest.fixture() diff --git a/tests/managers/thl/test_payout.py b/tests/managers/thl/test_payout.py index 0f3f103..2494de8 100644 --- a/tests/managers/thl/test_payout.py +++ b/tests/managers/thl/test_payout.py @@ -6,6 +6,7 @@ from datetime import UTC, datetime, timedelta from decimal import Decimal from random import choice as rand_choice from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pandas as pd @@ -13,35 +14,39 @@ import pytest from dask.distributed import Client as DaskClient from generalresearch.currency import USDCent -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections.thl_web import ( - LedgerDFCollection, -) -from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.payout import ( - BrokerageProductPayoutEventManager, - BusinessPayoutEventManager, - PayoutEventManager, - UserPayoutEventManager, -) -from generalresearch.managers.thl.product import ProductManager -from generalresearch.models.gr.business import Business from generalresearch.models.thl.definitions import PayoutStatus from generalresearch.models.thl.finance import BusinessBalances -from generalresearch.models.thl.ledger import LedgerAccount from generalresearch.models.thl.payout import ( BrokerageProductPayoutEvent, BusinessPayoutEvent, - UserPayoutEvent, ) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.models.thl.wallet import PayoutType -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig +from generalresearch.models.thl.wallet.definitions import PayoutType + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.thl_web import ( + LedgerDFCollection, + ) + from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.payout import ( + BrokerageProductPayoutEventManager, + BusinessPayoutEventManager, + PayoutEventManager, + UserPayoutEventManager, + ) + from generalresearch.managers.thl.product import ProductManager + from generalresearch.models.gr.business import Business + from generalresearch.models.thl.ledger import LedgerAccount + from generalresearch.models.thl.payout import ( + UserPayoutEvent, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig logger = logging.getLogger() @@ -189,7 +194,7 @@ class TestPayout: utc_now: datetime, ): from generalresearch.models.thl.definitions import PayoutStatus - from generalresearch.models.thl.wallet import PayoutType + from generalresearch.models.thl.wallet.definitions import PayoutType user_account = thl_ledger_manager.get_account_or_create_user_wallet(user=user) bp_account = thl_ledger_manager.get_account_or_create_bp_wallet(product=product) diff --git a/tests/managers/thl/test_product.py b/tests/managers/thl/test_product.py index f93ac36..644dc90 100644 --- a/tests/managers/thl/test_product.py +++ b/tests/managers/thl/test_product.py @@ -1,13 +1,12 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.product import ProductManager -from generalresearch.models import Source -from generalresearch.models.gr.team import Team +from generalresearch.models.definitions import Source from generalresearch.models.thl.product import ( Product, ProfilingConfig, @@ -19,6 +18,10 @@ from generalresearch.models.thl.product import ( UserHealthConfig, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.product import ProductManager + from generalresearch.models.gr.team import Team + class TestProductManagerGetMethods: def test_get_by_uuid(self, product_manager: ProductManager): diff --git a/tests/managers/thl/test_product_prod.py b/tests/managers/thl/test_product_prod.py index 8734210..d584527 100644 --- a/tests/managers/thl/test_product_prod.py +++ b/tests/managers/thl/test_product_prod.py @@ -2,13 +2,16 @@ from __future__ import annotations import logging from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.product import ProductManager from generalresearch.models.thl.product import Product +if TYPE_CHECKING: + from generalresearch.managers.thl.product import ProductManager + logger = logging.getLogger() diff --git a/tests/managers/thl/test_profiling/test_question.py b/tests/managers/thl/test_profiling/test_question.py index 97e7365..e4afb87 100644 --- a/tests/managers/thl/test_profiling/test_question.py +++ b/tests/managers/thl/test_profiling/test_question.py @@ -1,8 +1,11 @@ from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 -from generalresearch.managers.thl.profiling.question import QuestionManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source + +if TYPE_CHECKING: + from generalresearch.managers.thl.profiling.question import QuestionManager class TestQuestionManager: diff --git a/tests/managers/thl/test_profiling/test_schema.py b/tests/managers/thl/test_profiling/test_schema.py index b0eae31..feab902 100644 --- a/tests/managers/thl/test_profiling/test_schema.py +++ b/tests/managers/thl/test_profiling/test_schema.py @@ -1,10 +1,13 @@ from collections.abc import Callable +from typing import TYPE_CHECKING -from generalresearch.managers.thl.profiling.schema import ( - UpkSchemaManager, -) from generalresearch.models.thl.profiling.upk_property import PropertyType +if TYPE_CHECKING: + from generalresearch.managers.thl.profiling.schema import ( + UpkSchemaManager, + ) + class TestUpkSchemaManager: diff --git a/tests/managers/thl/test_profiling/test_user_upk.py b/tests/managers/thl/test_profiling/test_user_upk.py index fa10b67..0f3140c 100644 --- a/tests/managers/thl/test_profiling/test_user_upk.py +++ b/tests/managers/thl/test_profiling/test_user_upk.py @@ -1,8 +1,10 @@ from collections.abc import Callable from datetime import UTC, datetime +from typing import TYPE_CHECKING -from generalresearch.managers.thl.profiling.user_upk import UserUpkManager -from generalresearch.models.thl.user import User +if TYPE_CHECKING: + from generalresearch.managers.thl.profiling.user_upk import UserUpkManager + from generalresearch.models.thl.user import User now = datetime.now(tz=UTC) base = { diff --git a/tests/managers/thl/test_session_manager.py b/tests/managers/thl/test_session_manager.py index 05a49c1..30fd9ec 100644 --- a/tests/managers/thl/test_session_manager.py +++ b/tests/managers/thl/test_session_manager.py @@ -3,24 +3,27 @@ from __future__ import annotations from collections.abc import Callable from datetime import datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 from faker import Faker -from generalresearch.managers.thl.session import SessionManager -from generalresearch.models import DeviceType -from generalresearch.models.gr.business import Business -from generalresearch.models.gr.team import Team +from generalresearch.models.definitions import DeviceType from generalresearch.models.legacy.bucket import Bucket from generalresearch.models.thl.definitions import ( SessionStatusCode2, Status, StatusCode1, ) -from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import Session from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.managers.thl.session import SessionManager + from generalresearch.models.gr.business import Business + from generalresearch.models.gr.team import Team + from generalresearch.models.thl.product import Product + from generalresearch.pg_helper import PostgresConfig fake = Faker() diff --git a/tests/managers/thl/test_survey.py b/tests/managers/thl/test_survey.py index c3ab162..e114b70 100644 --- a/tests/managers/thl/test_survey.py +++ b/tests/managers/thl/test_survey.py @@ -4,16 +4,11 @@ import uuid from collections.abc import Callable from datetime import UTC, datetime from decimal import Decimal +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.buyer import BuyerManager -from generalresearch.managers.thl.profiling.question import ( - QuestionManager, -) -from generalresearch.managers.thl.profiling.uqa import UQAManager -from generalresearch.managers.thl.survey import SurveyManager, SurveyStatManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.legacy.bucket import ( DurationSummary, PayoutSummary, @@ -30,6 +25,14 @@ from generalresearch.models.thl.survey.model import ( SurveyStat, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.buyer import BuyerManager + from generalresearch.managers.thl.profiling.question import ( + QuestionManager, + ) + from generalresearch.managers.thl.profiling.uqa import UQAManager + from generalresearch.managers.thl.survey import SurveyManager, SurveyStatManager + @pytest.fixture(scope="session") def surveys_fixture() -> list[Survey]: @@ -223,7 +226,6 @@ class TestSurvey: class TestSurveyStat: def test( self, - delete_buyers_surveys: Callable[..., None], surveystat_manager, survey_manager: SurveyManager, surveys_fixture: list[Survey], diff --git a/tests/managers/thl/test_survey_penalty.py b/tests/managers/thl/test_survey_penalty.py index 9c29a0a..04f69d2 100644 --- a/tests/managers/thl/test_survey_penalty.py +++ b/tests/managers/thl/test_survey_penalty.py @@ -1,16 +1,19 @@ from __future__ import annotations import uuid +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.survey_penalty import SurveyPenaltyManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.survey.penalty import ( BPSurveyPenalty, TeamSurveyPenalty, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.survey_penalty import SurveyPenaltyManager + @pytest.fixture def product_uuid() -> str: diff --git a/tests/managers/thl/test_task_adjustment.py b/tests/managers/thl/test_task_adjustment.py index a7324c3..a14401e 100644 --- a/tests/managers/thl/test_task_adjustment.py +++ b/tests/managers/thl/test_task_adjustment.py @@ -5,23 +5,26 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal from random import randint +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.task_adjustment import ( - TaskAdjustmentManager, -) -from generalresearch.managers.thl.wall import WallManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( Status, StatusCode1, WallAdjustedStatus, ) -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.task_adjustment import ( + TaskAdjustmentManager, + ) + from generalresearch.managers.thl.wall import WallManager + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User @pytest.fixture() diff --git a/tests/managers/thl/test_task_status.py b/tests/managers/thl/test_task_status.py index b47f650..9846ce0 100644 --- a/tests/managers/thl/test_task_status.py +++ b/tests/managers/thl/test_task_status.py @@ -3,13 +3,11 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.product import ProductManager -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.wall import WallManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( Status, StatusCode1, @@ -19,12 +17,18 @@ from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, PayoutTransformationPercentArgs, - Product, UserWalletConfig, ) -from generalresearch.models.thl.session import Session, WallOut +from generalresearch.models.thl.session import WallOut from generalresearch.models.thl.task_status import TaskStatusResponse -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.product import ProductManager + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.wall import WallManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User start1 = datetime(2023, 2, 1, tzinfo=UTC) finish1 = start1 + timedelta(minutes=5) diff --git a/tests/managers/thl/test_user_manager/test_base.py b/tests/managers/thl/test_user_manager/test_base.py index 8cd83ad..4a9750e 100644 --- a/tests/managers/thl/test_user_manager/test_base.py +++ b/tests/managers/thl/test_user_manager/test_base.py @@ -1,11 +1,11 @@ import logging from datetime import UTC, datetime from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.product import ProductManager from generalresearch.managers.thl.user_manager import ( UserCreateNotAllowedError, get_bp_user_create_limit_hourly, @@ -17,13 +17,17 @@ from generalresearch.managers.thl.user_manager.rate_limit import ( RateLimitItemPerHourConstantKey, UserManagerLimiter, ) -from generalresearch.managers.thl.user_manager.user_manager import ( - UserManager, -) -from generalresearch.managers.thl.userhealth import AuditLogManager -from generalresearch.models.thl.product import Product, UserCreateConfig +from generalresearch.models.thl.product import UserCreateConfig from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.managers.thl.product import ProductManager + from generalresearch.managers.thl.user_manager.user_manager import ( + UserManager, + ) + from generalresearch.managers.thl.userhealth import AuditLogManager + from generalresearch.models.thl.product import Product + from generalresearch.pg_helper import PostgresConfig logger = logging.getLogger() diff --git a/tests/managers/thl/test_user_manager/test_mysql.py b/tests/managers/thl/test_user_manager/test_mysql.py index e6f43ef..ed7d458 100644 --- a/tests/managers/thl/test_user_manager/test_mysql.py +++ b/tests/managers/thl/test_user_manager/test_mysql.py @@ -1,9 +1,12 @@ from __future__ import annotations -from generalresearch.managers.thl.user_manager.mysql_user_manager import ( - MysqlUserManager, -) -from generalresearch.models.thl.user import User +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from generalresearch.managers.thl.user_manager.mysql_user_manager import ( + MysqlUserManager, + ) + from generalresearch.models.thl.user import User class TestUserManagerMysqlNew: diff --git a/tests/managers/thl/test_user_manager/test_redis.py b/tests/managers/thl/test_user_manager/test_redis.py index 04071ee..e51aae9 100644 --- a/tests/managers/thl/test_user_manager/test_redis.py +++ b/tests/managers/thl/test_user_manager/test_redis.py @@ -1,14 +1,18 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest -from generalresearch.config import GRLBaseSettings from generalresearch.managers.base import Permission from generalresearch.managers.thl.user_manager.redis_user_manager import ( RedisUserManager, ) -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig + +if TYPE_CHECKING: + from generalresearch.config import GRLBaseSettings + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig class TestUserManagerRedis: diff --git a/tests/managers/thl/test_user_manager/test_user_fetch.py b/tests/managers/thl/test_user_manager/test_user_fetch.py index 87d010a..9a279ed 100644 --- a/tests/managers/thl/test_user_manager/test_user_fetch.py +++ b/tests/managers/thl/test_user_manager/test_user_fetch.py @@ -1,13 +1,15 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.user_manager.user_manager import UserManager -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User +if TYPE_CHECKING: + from generalresearch.managers.thl.user_manager.user_manager import UserManager + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User class TestUserManagerFetch: diff --git a/tests/managers/thl/test_user_manager/test_user_metadata.py b/tests/managers/thl/test_user_manager/test_user_metadata.py index 670e38a..eb6a272 100644 --- a/tests/managers/thl/test_user_manager/test_user_metadata.py +++ b/tests/managers/thl/test_user_manager/test_user_metadata.py @@ -1,17 +1,20 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.user_manager.user_metadata_manager import ( - UserMetadataManager, -) -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User from generalresearch.models.thl.user_profile import UserMetadata +if TYPE_CHECKING: + from generalresearch.managers.thl.user_manager.user_metadata_manager import ( + UserMetadataManager, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + class TestUserMetadataManager: diff --git a/tests/managers/thl/test_user_streak.py b/tests/managers/thl/test_user_streak.py index 61e2947..564a142 100644 --- a/tests/managers/thl/test_user_streak.py +++ b/tests/managers/thl/test_user_streak.py @@ -3,17 +3,15 @@ from __future__ import annotations import copy from datetime import UTC, date, datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING from zoneinfo import ZoneInfo import pytest -from generalresearch.managers.thl.session import SessionManager from generalresearch.managers.thl.user_streak import ( - UserStreakManager, compute_streaks_from_days, ) from generalresearch.models.thl.definitions import Status, StatusCode1 -from generalresearch.models.thl.user import User from generalresearch.models.thl.user_streak import ( StreakFulfillment, StreakPeriod, @@ -21,6 +19,13 @@ from generalresearch.models.thl.user_streak import ( UserStreak, ) +if TYPE_CHECKING: + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.user_streak import ( + UserStreakManager, + ) + from generalresearch.models.thl.user import User + def test_compute_streaks_from_days(): days = [ diff --git a/tests/managers/thl/test_userhealth.py b/tests/managers/thl/test_userhealth.py index ea54359..ce6c221 100644 --- a/tests/managers/thl/test_userhealth.py +++ b/tests/managers/thl/test_userhealth.py @@ -2,6 +2,7 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime +from typing import TYPE_CHECKING from uuid import uuid4 import faker @@ -12,16 +13,24 @@ from generalresearch.managers.thl.userhealth import ( IPRecordManager, UserIpHistoryManager, ) -from generalresearch.models.thl.ipinfo import GeoIPInformation, IPGeoname, IPInformation -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User +from generalresearch.models.thl.ipinfo import ( + GeoIPInformation, +) from generalresearch.models.thl.user_iphistory import ( IPRecord, UserIPHistory, ) from generalresearch.models.thl.userhealth import AuditLog, AuditLogLevel -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.models.thl.ipinfo import ( + IPGeoname, + IPInformation, + ) + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig fake = faker.Faker() diff --git a/tests/managers/thl/test_wall_manager.py b/tests/managers/thl/test_wall_manager.py index b8a636f..3215de8 100644 --- a/tests/managers/thl/test_wall_manager.py +++ b/tests/managers/thl/test_wall_manager.py @@ -3,21 +3,24 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pydantic import PositiveInt -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.wall import WallCacheManager, WallManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.session import ( ReportValue, - Session, Status, StatusCode1, ) -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.wall import WallCacheManager, WallManager + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User class TestWallManager: diff --git a/tests/models/custom_types/test_aware_datetime.py b/tests/models/custom_types/test_aware_datetime.py index e8a5aa3..54a5d9b 100644 --- a/tests/models/custom_types/test_aware_datetime.py +++ b/tests/models/custom_types/test_aware_datetime.py @@ -2,12 +2,14 @@ from __future__ import annotations import logging from datetime import UTC, datetime +from typing import TYPE_CHECKING import pytest import pytz from pydantic import BaseModel, Field, ValidationError -from generalresearch.models.custom_types import AwareDatetimeISO +if TYPE_CHECKING: + from generalresearch.models.custom_types import AwareDatetimeISO logger = logging.getLogger() diff --git a/tests/models/custom_types/test_dsn.py b/tests/models/custom_types/test_dsn.py index d8c7c53..2aae579 100644 --- a/tests/models/custom_types/test_dsn.py +++ b/tests/models/custom_types/test_dsn.py @@ -1,12 +1,14 @@ from __future__ import annotations +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pydantic import BaseModel, Field, MySQLDsn, ValidationError from pydantic_core import Url -from generalresearch.models.custom_types import DaskDsn, SentryDsn +if TYPE_CHECKING: + from generalresearch.models.custom_types import DaskDsn, SentryDsn # --- Test Pydantic Models --- diff --git a/tests/models/custom_types/test_uuid_str.py b/tests/models/custom_types/test_uuid_str.py index 02e6a8b..92489a0 100644 --- a/tests/models/custom_types/test_uuid_str.py +++ b/tests/models/custom_types/test_uuid_str.py @@ -1,11 +1,13 @@ from __future__ import annotations +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pydantic import BaseModel, Field, ValidationError -from generalresearch.models.custom_types import UUIDStr +if TYPE_CHECKING: + from generalresearch.models.custom_types import UUIDStr class UUIDStrModel(BaseModel): diff --git a/tests/models/dynata/test_survey.py b/tests/models/dynata/test_survey.py index ad953a3..3e33897 100644 --- a/tests/models/dynata/test_survey.py +++ b/tests/models/dynata/test_survey.py @@ -1,3 +1,6 @@ +from __future__ import annotations + + class TestDynataCondition: def test_condition_create(self): diff --git a/tests/models/gr/test_authentication.py b/tests/models/gr/test_authentication.py index d2a7054..ac1298f 100644 --- a/tests/models/gr/test_authentication.py +++ b/tests/models/gr/test_authentication.py @@ -6,17 +6,21 @@ import os from collections.abc import Callable from datetime import UTC, datetime from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from redis import Redis from generalresearch.models.gr.authentication import Claims, GRToken, GRUser -from generalresearch.models.gr.business import Business -from generalresearch.models.gr.team import Membership, Team -from generalresearch.models.thl.product import Product -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig +from generalresearch.models.gr.team import Team + +if TYPE_CHECKING: + from generalresearch.models.gr.business import Business + from generalresearch.models.gr.team import Membership + from generalresearch.models.thl.product import Product + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig SSO_ISSUER = "" diff --git a/tests/models/gr/test_base.py b/tests/models/gr/test_base.py index f019fc0..fba0960 100644 --- a/tests/models/gr/test_base.py +++ b/tests/models/gr/test_base.py @@ -3,11 +3,13 @@ from __future__ import annotations import subprocess from collections.abc import Callable from pathlib import Path +from typing import TYPE_CHECKING import pytest from pydantic import PostgresDsn -from generalresearch.pg_helper import PostgresConfig +if TYPE_CHECKING: + from generalresearch.pg_helper import PostgresConfig class TestGRPostgresDjangoCreation: diff --git a/tests/models/gr/test_business.py b/tests/models/gr/test_business.py index 9310d2c..2c12da1 100644 --- a/tests/models/gr/test_business.py +++ b/tests/models/gr/test_business.py @@ -19,40 +19,42 @@ from distributed.utils_test import ( from pytest import approx from generalresearch.currency import USDCent -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - WallDFCollection, -) -from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge -from generalresearch.managers.gr.business import BusinessBankAccountManager -from generalresearch.managers.gr.team import TeamManager -from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.payout import ( - BusinessPayoutEventManager, - PayoutEventManager, -) from generalresearch.models.gr.business import ( Business, BusinessAddress, - BusinessBankAccount, BusinessContact, ) -from generalresearch.models.gr.team import Team from generalresearch.models.thl.finance import ( BusinessBalances, ProductBalances, ) from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + WallDFCollection, + ) + from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge + from generalresearch.managers.gr.business import BusinessBankAccountManager + from generalresearch.managers.gr.team import TeamManager + from generalresearch.managers.thl.ledger_manager.ledger import LedgerManager + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.payout import ( + BusinessPayoutEventManager, + PayoutEventManager, + ) from generalresearch.managers.thl.product import ProductManager + from generalresearch.models.gr.business import ( + BusinessBankAccount, + ) + from generalresearch.models.gr.team import Team from generalresearch.models.thl.product import BrokerageProductPayoutEvent from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig class TestBusinessBankAccount: diff --git a/tests/models/gr/test_team.py b/tests/models/gr/test_team.py index dc7d4b9..c1ae6d6 100644 --- a/tests/models/gr/test_team.py +++ b/tests/models/gr/test_team.py @@ -5,6 +5,7 @@ from collections.abc import Callable from datetime import datetime, timedelta from decimal import Decimal from pathlib import Path +from typing import TYPE_CHECKING import pandas as pd from dask.distributed import Client as DaskClient @@ -12,24 +13,29 @@ from distributed.utils_test import ( client_no_amm, ) -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections.thl_web import ( - SessionDFCollection, - WallDFCollection, -) -from generalresearch.incite.mergers.foundations.enriched_session import ( - EnrichedSessionMerge, -) -from generalresearch.incite.mergers.foundations.enriched_wall import EnrichedWallMerge -from generalresearch.managers.gr.team import MembershipManager, TeamManager -from generalresearch.models.gr.authentication import GRUser from generalresearch.models.gr.business import Business -from generalresearch.models.gr.team import Membership, Team +from generalresearch.models.gr.team import Team from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.thl_web import ( + SessionDFCollection, + WallDFCollection, + ) + from generalresearch.incite.mergers.foundations.enriched_session import ( + EnrichedSessionMerge, + ) + from generalresearch.incite.mergers.foundations.enriched_wall import ( + EnrichedWallMerge, + ) + from generalresearch.managers.gr.team import MembershipManager, TeamManager + from generalresearch.models.gr.authentication import GRUser + from generalresearch.models.gr.team import Membership + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + from generalresearch.redis_helper import RedisConfig class TestTeam: diff --git a/tests/models/innovate/test_question.py b/tests/models/innovate/test_question.py index b206177..ea2fc8c 100644 --- a/tests/models/innovate/test_question.py +++ b/tests/models/innovate/test_question.py @@ -1,6 +1,6 @@ from __future__ import annotations -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.innovate.question import ( InnovateQuestion, InnovateQuestionOption, diff --git a/tests/models/legacy/test_offerwall_parse_response.py b/tests/models/legacy/test_offerwall_parse_response.py index 56ba077..93f5c26 100644 --- a/tests/models/legacy/test_offerwall_parse_response.py +++ b/tests/models/legacy/test_offerwall_parse_response.py @@ -2,7 +2,7 @@ from __future__ import annotations import json -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.legacy.bucket import ( BucketTask, DurationSummary, diff --git a/tests/models/legacy/test_user_question_answer_in.py b/tests/models/legacy/test_user_question_answer_in.py index 3fdaa05..f14c1a7 100644 --- a/tests/models/legacy/test_user_question_answer_in.py +++ b/tests/models/legacy/test_user_question_answer_in.py @@ -4,19 +4,22 @@ import json from collections.abc import Callable from datetime import datetime from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest -from generalresearch.managers.thl.user_manager.user_manager import UserManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.legacy.questions import ( UserQuestionAnswers, ) -from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import Session, Wall from generalresearch.models.thl.user import User +if TYPE_CHECKING: + from generalresearch.managers.thl.user_manager.user_manager import UserManager + from generalresearch.models.thl.product import Product + class TestUserQuestionAnswers: """This is for the GRS POST submission that may contain multiple diff --git a/tests/models/network/test_mtr.py b/tests/models/network/test_mtr.py index 7f8a736..5d136c4 100644 --- a/tests/models/network/test_mtr.py +++ b/tests/models/network/test_mtr.py @@ -1,11 +1,15 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import faker -from generalresearch.managers.network.tool_run import ToolRunManager from generalresearch.models.network.mtr.execute import execute_mtr from generalresearch.models.network.tool_run import ToolClass, ToolName +if TYPE_CHECKING: + from generalresearch.managers.network.tool_run import ToolRunManager + fake = faker.Faker() diff --git a/tests/models/network/test_nmap.py b/tests/models/network/test_nmap.py index db39997..6adc9e4 100644 --- a/tests/models/network/test_nmap.py +++ b/tests/models/network/test_nmap.py @@ -1,14 +1,18 @@ from __future__ import annotations import subprocess +from typing import TYPE_CHECKING import faker -from generalresearch.managers.network.tool_run import ToolRunManager from generalresearch.models.network.definitions import IPProtocol from generalresearch.models.network.nmap.execute import execute_nmap from generalresearch.models.network.nmap.result import NmapResult, PortState -from generalresearch.models.network.tool_run import NmapRun, ToolClass, ToolName +from generalresearch.models.network.tool_run import ToolClass, ToolName + +if TYPE_CHECKING: + from generalresearch.managers.network.tool_run import ToolRunManager + from generalresearch.models.network.tool_run import NmapRun fake = faker.Faker() diff --git a/tests/models/network/test_nmap_parser.py b/tests/models/network/test_nmap_parser.py index 473a63f..fc9884b 100644 --- a/tests/models/network/test_nmap_parser.py +++ b/tests/models/network/test_nmap_parser.py @@ -1,14 +1,15 @@ from __future__ import annotations import os +from typing import TYPE_CHECKING import pytest from generalresearch.models.network.nmap.parser import parse_nmap_xml -from generalresearch.models.network.nmap.result import ( - NmapResult, - NmapTrace, -) +from generalresearch.models.network.nmap.result import NmapTrace + +if TYPE_CHECKING: + from generalresearch.models.network.nmap.result import NmapResult @pytest.fixture diff --git a/tests/models/network/test_rdns.py b/tests/models/network/test_rdns.py index 1a15a28..82126dd 100644 --- a/tests/models/network/test_rdns.py +++ b/tests/models/network/test_rdns.py @@ -1,11 +1,15 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import faker -from generalresearch.managers.network.tool_run import ToolRunManager from generalresearch.models.network.rdns.execute import execute_rdns from generalresearch.models.network.tool_run import ToolClass, ToolName +if TYPE_CHECKING: + from generalresearch.managers.network.tool_run import ToolRunManager + fake = faker.Faker() diff --git a/tests/models/spectrum/test_question.py b/tests/models/spectrum/test_question.py index a44286d..d469530 100644 --- a/tests/models/spectrum/test_question.py +++ b/tests/models/spectrum/test_question.py @@ -2,7 +2,7 @@ from __future__ import annotations from datetime import UTC, datetime -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.spectrum.question import ( SpectrumQuestion, SpectrumQuestionClass, diff --git a/tests/models/spectrum/test_survey.py b/tests/models/spectrum/test_survey.py index bad6857..02c5d3f 100644 --- a/tests/models/spectrum/test_survey.py +++ b/tests/models/spectrum/test_survey.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import UTC, datetime from decimal import Decimal -from generalresearch.models import ( +from generalresearch.models.definitions import ( LogicalOperator, Source, TaskCalculationType, diff --git a/tests/models/spectrum/test_survey_manager.py b/tests/models/spectrum/test_survey_manager.py index 11dc01f..0300956 100644 --- a/tests/models/spectrum/test_survey_manager.py +++ b/tests/models/spectrum/test_survey_manager.py @@ -3,15 +3,17 @@ from __future__ import annotations import logging from datetime import UTC, datetime from decimal import Decimal -from typing import Any +from typing import TYPE_CHECKING, Any from pymysql import IntegrityError from generalresearch.config import is_debug -from generalresearch.managers.spectrum.survey import ( - SpectrumSurveyManager, -) -from generalresearch.sql_helper import SqlHelper + +if TYPE_CHECKING: + from generalresearch.managers.spectrum.survey import ( + SpectrumSurveyManager, + ) + from generalresearch.sql_helper import SqlHelper logger = logging.getLogger() diff --git a/tests/models/test_device.py b/tests/models/test_device.py index 8e1251a..fdbd906 100644 --- a/tests/models/test_device.py +++ b/tests/models/test_device.py @@ -15,7 +15,7 @@ chromebook_ua_string = ( ) -from generalresearch.models import DeviceType +from generalresearch.models.definitions import DeviceType from generalresearch.models.device import parse_device_from_useragent diff --git a/tests/models/test_finance.py b/tests/models/test_finance.py index 72f4f4d..eabc877 100644 --- a/tests/models/test_finance.py +++ b/tests/models/test_finance.py @@ -4,6 +4,7 @@ from collections.abc import Callable from datetime import UTC, datetime, timedelta from itertools import product as iter_product from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pandas as pd @@ -16,25 +17,27 @@ from distributed.utils_test import ( ) from faker import Faker -from generalresearch.incite.collections.thl_web import LedgerDFCollection -from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge from generalresearch.incite.schemas.mergers.pop_ledger import ( numerical_col_names, ) -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager from generalresearch.models.thl.finance import ( BusinessBalances, POPFinancial, ProductBalances, ) -from generalresearch.models.thl.ledger import LedgerAccount -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.pg_helper import PostgresConfig from test_utils.incite.collections.conftest import ledger_collection from test_utils.incite.mergers.conftest import pop_ledger_merge +if TYPE_CHECKING: + from generalresearch.incite.collections.thl_web import LedgerDFCollection + from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.models.thl.ledger import LedgerAccount + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.pg_helper import PostgresConfig + fake = Faker() diff --git a/tests/models/thl/test_adjustments.py b/tests/models/thl/test_adjustments.py index c5b3f6b..cd75318 100644 --- a/tests/models/thl/test_adjustments.py +++ b/tests/models/thl/test_adjustments.py @@ -3,22 +3,27 @@ from __future__ import annotations from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING import pytest -from generalresearch.managers.thl.session import SessionManager -from generalresearch.managers.thl.wall import WallManager -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.product import Product from generalresearch.models.thl.session import ( - Session, SessionAdjustedStatus, Status, StatusCode1, - Wall, WallAdjustedStatus, ) -from generalresearch.models.thl.user import User + +if TYPE_CHECKING: + from generalresearch.managers.thl.session import SessionManager + from generalresearch.managers.thl.wall import WallManager + from generalresearch.models.thl.session import ( + Session, + Wall, + ) + from generalresearch.models.thl.user import User started1 = datetime(2023, 1, 1, tzinfo=UTC) started2 = datetime(2023, 1, 1, 0, 10, 0, tzinfo=UTC) diff --git a/tests/models/thl/test_buyer.py b/tests/models/thl/test_buyer.py index 02093e2..ef97166 100644 --- a/tests/models/thl/test_buyer.py +++ b/tests/models/thl/test_buyer.py @@ -1,6 +1,6 @@ from __future__ import annotations -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.survey.buyer import BuyerCountryStat diff --git a/tests/models/thl/test_contest/test_contest.py b/tests/models/thl/test_contest/test_contest.py index e1053f4..ed8477b 100644 --- a/tests/models/thl/test_contest/test_contest.py +++ b/tests/models/thl/test_contest/test_contest.py @@ -1,11 +1,13 @@ from __future__ import annotations from collections.abc import Callable +from typing import TYPE_CHECKING import pytest -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User +if TYPE_CHECKING: + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User class TestContest: diff --git a/tests/models/thl/test_contest/test_leaderboard_contest.py b/tests/models/thl/test_contest/test_leaderboard_contest.py index 99cfb37..c49776b 100644 --- a/tests/models/thl/test_contest/test_leaderboard_contest.py +++ b/tests/models/thl/test_contest/test_leaderboard_contest.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -8,7 +9,6 @@ from redis import Redis from generalresearch.currency import USDCent from generalresearch.managers.leaderboard.manager import LeaderboardManager -from generalresearch.managers.thl.user_manager.user_manager import UserManager from generalresearch.models.thl.contest import ContestPrize from generalresearch.models.thl.contest.definitions import ( ContestPrizeKind, @@ -21,10 +21,13 @@ from generalresearch.models.thl.contest.utils import ( distribute_leaderboard_prizes, ) from generalresearch.models.thl.leaderboard import LeaderboardRow -from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User from tests.models.thl.test_contest.test_contest import TestContest +if TYPE_CHECKING: + from generalresearch.managers.thl.user_manager.user_manager import UserManager + from generalresearch.models.thl.product import Product + class TestLeaderboardContest(TestContest): diff --git a/tests/models/thl/test_contest/test_raffle_contest.py b/tests/models/thl/test_contest/test_raffle_contest.py index 8812cb3..e71851e 100644 --- a/tests/models/thl/test_contest/test_raffle_contest.py +++ b/tests/models/thl/test_contest/test_raffle_contest.py @@ -2,6 +2,7 @@ from __future__ import annotations from collections import Counter from datetime import datetime +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -21,10 +22,12 @@ from generalresearch.models.thl.contest.definitions import ( ContestType, ) from generalresearch.models.thl.contest.raffle import RaffleContest -from generalresearch.models.thl.product import Product -from generalresearch.models.thl.user import User from tests.models.thl.test_contest.test_contest import TestContest +if TYPE_CHECKING: + from generalresearch.models.thl.product import Product + from generalresearch.models.thl.user import User + class TestRaffleContest(TestContest): diff --git a/tests/models/thl/test_marketplace_condition.py b/tests/models/thl/test_marketplace_condition.py index 1dd25e8..6936a7c 100644 --- a/tests/models/thl/test_marketplace_condition.py +++ b/tests/models/thl/test_marketplace_condition.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest from pydantic import ValidationError -from generalresearch.models import LogicalOperator +from generalresearch.models.definitions import LogicalOperator from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, @@ -130,7 +130,7 @@ class TestMarketplaceCondition: assert c.evaluate_criterion(user_qas) is None def test_list_and_negate(self): - from generalresearch.models import LogicalOperator + from generalresearch.models.definitions import LogicalOperator from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, @@ -247,7 +247,7 @@ class TestMarketplaceCondition: assert ["1", "10", "11", "12", "2", "3", "4", "5"] == c.values def test_ranges_infinity(self): - from generalresearch.models import LogicalOperator + from generalresearch.models.definitions import LogicalOperator from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, diff --git a/tests/models/thl/test_payout.py b/tests/models/thl/test_payout.py index daf1bd7..927687e 100644 --- a/tests/models/thl/test_payout.py +++ b/tests/models/thl/test_payout.py @@ -7,12 +7,16 @@ from pydantic import ValidationError from generalresearch.currency import USDCent from generalresearch.models.gr import Team -from generalresearch.models.gr.business import Business, BusinessAddress, BusinessType +from generalresearch.models.gr.business import ( + Business, + BusinessAddress, + BusinessType, +) from generalresearch.models.thl.payout import ( BrokerageProductPayoutEvent, BusinessPayoutEvent, ) -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType class TestBusinessPayoutEvent: diff --git a/tests/models/thl/test_payout_format.py b/tests/models/thl/test_payout_format.py index fe7aea5..56eafe3 100644 --- a/tests/models/thl/test_payout_format.py +++ b/tests/models/thl/test_payout_format.py @@ -1,14 +1,20 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import pytest from pydantic import BaseModel from generalresearch.models.thl.payout_format import ( PayoutFormatField, - PayoutFormatType, format_payout_format, ) +if TYPE_CHECKING: + from generalresearch.models.thl.payout_format import ( + PayoutFormatType, + ) + class PayoutFormatTestClass(BaseModel): payout_format: PayoutFormatType = PayoutFormatField diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index 880799a..cc0fa8e 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -5,6 +5,7 @@ import shutil from collections.abc import Callable from datetime import UTC, datetime, timedelta from decimal import Decimal +from typing import TYPE_CHECKING from uuid import uuid4 import pytest @@ -12,22 +13,9 @@ from dask.distributed import Client as DaskClient from pydantic import ValidationError from generalresearch.currency import USDCent -from generalresearch.incite.base import GRLDatasets -from generalresearch.incite.collections.thl_web import LedgerDFCollection -from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge -from generalresearch.managers.thl.ledger_manager.thl_ledger import ( - ThlLedgerManager, -) -from generalresearch.managers.thl.payout import PayoutEventManager -from generalresearch.managers.thl.product import ProductManager -from generalresearch.models import Source -from generalresearch.models.gr.business import Business +from generalresearch.models.definitions import Source from generalresearch.models.thl.finance import ProductBalances -from generalresearch.models.thl.payout import ( - BrokerageProductPayoutEvent, -) from generalresearch.models.thl.product import ( - BrokerageProductPayoutEventManager, IntegrationMode, PayoutConfig, PayoutTransformation, @@ -39,9 +27,24 @@ from generalresearch.models.thl.product import ( SupplyConfig, SupplyPolicy, ) -from generalresearch.models.thl.session import Session -from generalresearch.models.thl.user import User -from generalresearch.redis_helper import RedisConfig + +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + from generalresearch.incite.collections.thl_web import LedgerDFCollection + from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge + from generalresearch.managers.thl.ledger_manager.thl_ledger import ( + ThlLedgerManager, + ) + from generalresearch.managers.thl.payout import PayoutEventManager + from generalresearch.managers.thl.product import ProductManager + from generalresearch.models.gr.business import Business + from generalresearch.models.thl.payout import ( + BrokerageProductPayoutEvent, + ) + from generalresearch.models.thl.product import BrokerageProductPayoutEventManager + from generalresearch.models.thl.session import Session + from generalresearch.models.thl.user import User + from generalresearch.redis_helper import RedisConfig class TestProduct: diff --git a/tests/models/thl/test_product_userwalletconfig.py b/tests/models/thl/test_product_userwalletconfig.py index b348981..b1bfa60 100644 --- a/tests/models/thl/test_product_userwalletconfig.py +++ b/tests/models/thl/test_product_userwalletconfig.py @@ -6,7 +6,7 @@ from random import shuffle as rshuffle from generalresearch.models.thl.product import ( UserWalletConfig, ) -from generalresearch.models.thl.wallet import PayoutType +from generalresearch.models.thl.wallet.definitions import PayoutType def all_equal(iterable: list[str]) -> bool: diff --git a/tests/models/thl/test_soft_pair.py b/tests/models/thl/test_soft_pair.py index 3cf835e..34902e2 100644 --- a/tests/models/thl/test_soft_pair.py +++ b/tests/models/thl/test_soft_pair.py @@ -1,6 +1,6 @@ from __future__ import annotations -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.dynata.survey import ( ConditionValueType, DynataCondition, diff --git a/tests/models/thl/test_user.py b/tests/models/thl/test_user.py index 9c4b548..bc941d4 100644 --- a/tests/models/thl/test_user.py +++ b/tests/models/thl/test_user.py @@ -6,16 +6,19 @@ from datetime import UTC, datetime, timedelta, timezone from decimal import Decimal from random import choice as rand_choice from random import randint +from typing import TYPE_CHECKING from uuid import uuid4 import pytest from pydantic import ValidationError -from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager -from generalresearch.managers.thl.userhealth import AuditLogManager -from generalresearch.models.thl.product import Product from generalresearch.models.thl.user import User +if TYPE_CHECKING: + from generalresearch.managers.thl.ledger_manager.thl_ledger import ThlLedgerManager + from generalresearch.managers.thl.userhealth import AuditLogManager + from generalresearch.models.thl.product import Product + class TestUserUserID: diff --git a/tests/models/thl/test_user_metadata.py b/tests/models/thl/test_user_metadata.py index a7b479d..7e84f3e 100644 --- a/tests/models/thl/test_user_metadata.py +++ b/tests/models/thl/test_user_metadata.py @@ -2,7 +2,7 @@ from __future__ import annotations import pytest -from generalresearch.models import MAX_INT32 +from generalresearch.models.definitions import MAX_INT32 from generalresearch.models.thl.user_profile import UserMetadata diff --git a/tests/models/thl/test_wall.py b/tests/models/thl/test_wall.py index 58e9825..61ca11d 100644 --- a/tests/models/thl/test_wall.py +++ b/tests/models/thl/test_wall.py @@ -7,7 +7,7 @@ from uuid import uuid4 import pytest from pydantic import ValidationError -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import ( Status, StatusCode1, diff --git a/tests/models/thl/test_wall_session.py b/tests/models/thl/test_wall_session.py index 48b89ea..40d3619 100644 --- a/tests/models/thl/test_wall_session.py +++ b/tests/models/thl/test_wall_session.py @@ -5,7 +5,7 @@ from decimal import Decimal import pytest -from generalresearch.models import Source +from generalresearch.models.definitions import Source from generalresearch.models.thl.definitions import Status, StatusCode1 from generalresearch.models.thl.session import Session, Wall from generalresearch.models.thl.user import User diff --git a/tests/test_postgres.py b/tests/test_postgres.py index ed5a7ae..c53f644 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -1,12 +1,15 @@ import socket import subprocess from collections.abc import Callable +from typing import TYPE_CHECKING from pydantic import PostgresDsn -from generalresearch.models.custom_types import InternalHostname, PostgresDict from generalresearch.pg_helper import PostgresConfig +if TYPE_CHECKING: + from generalresearch.models.custom_types import InternalHostname, PostgresDict + def is_port_open(host: InternalHostname, port: int = 5432, timeout: int = 3): try: -- cgit v1.2.3 From 48863e9431d50fd86405036b55b333b7734b00ac Mon Sep 17 00:00:00 2001 From: Max Nanis Date: Thu, 3 Sep 2026 11:30:55 -0700 Subject: import fixes for <3.13 --- test_utils/models/network/conftest.py | 4 +--- test_utils/models/thl/conftest.py | 2 +- test_utils/models/upk/conftest.py | 6 ++---- 3 files changed, 4 insertions(+), 8 deletions(-) (limited to 'test_utils/models/network') diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py index 4ff59ee..c7bcc7e 100644 --- a/test_utils/models/network/conftest.py +++ b/test_utils/models/network/conftest.py @@ -24,9 +24,7 @@ from generalresearch.models.network.tool_run_command import ( RDNSRunCommand, RDNSRunCommandOptions, ) - -if TYPE_CHECKING: - from generalresearch.pg_helper import PostgresConfig +from generalresearch.pg_helper import PostgresConfig @pytest.fixture(scope="session") diff --git a/test_utils/models/thl/conftest.py b/test_utils/models/thl/conftest.py index 021b19e..f1cb785 100644 --- a/test_utils/models/thl/conftest.py +++ b/test_utils/models/thl/conftest.py @@ -32,6 +32,7 @@ from generalresearch.models.thl.user import User from generalresearch.models.thl.user_iphistory import IPRecord from generalresearch.models.thl.userhealth import AuditLogLevel from generalresearch.models.thl.wallet.definitions import PayoutType +from generalresearch.pg_helper import PostgresConfig if TYPE_CHECKING: from generalresearch.currency import USDCent @@ -73,7 +74,6 @@ if TYPE_CHECKING: from generalresearch.models.thl.user_iphistory import IPRecord from generalresearch.models.thl.userhealth import AuditLog from generalresearch.models.thl.wallet.cashout_method import CashMailOrderData - from generalresearch.pg_helper import PostgresConfig fake = faker.Faker() diff --git a/test_utils/models/upk/conftest.py b/test_utils/models/upk/conftest.py index 59266b2..ad96bbb 100644 --- a/test_utils/models/upk/conftest.py +++ b/test_utils/models/upk/conftest.py @@ -3,15 +3,13 @@ from __future__ import annotations import os import time from collections.abc import Callable -from typing import TYPE_CHECKING from uuid import UUID import pandas as pd import pytest -if TYPE_CHECKING: - from generalresearch.managers.thl.category import CategoryManager - from generalresearch.pg_helper import PostgresConfig +from generalresearch.managers.thl.category import CategoryManager +from generalresearch.pg_helper import PostgresConfig def insert_data_from_csv( -- cgit v1.2.3 From 4b705346968e38671ac9601cfa8444c944ecc8bd Mon Sep 17 00:00:00 2001 From: stuppie Date: Fri, 4 Sep 2026 12:52:46 -0600 Subject: remove a lot of if type_checking; fix model_rebuild errors. Remove all network models/managers/test/fixtures --- generalresearch/managers/network/__init__.py | 0 generalresearch/managers/network/label.py | 152 ------- generalresearch/managers/network/mtr.py | 53 --- generalresearch/managers/network/nmap.py | 60 --- generalresearch/managers/network/rdns.py | 37 -- generalresearch/managers/network/tool_run.py | 138 ------- generalresearch/models/cint/question.py | 12 +- generalresearch/models/cint/survey.py | 13 +- generalresearch/models/cint/task_collection.py | 6 +- generalresearch/models/dynata/survey.py | 10 +- generalresearch/models/dynata/task_collection.py | 6 +- generalresearch/models/innovate/question.py | 12 +- generalresearch/models/innovate/survey.py | 10 +- generalresearch/models/innovate/task_collection.py | 6 +- generalresearch/models/legacy/offerwall.py | 37 +- generalresearch/models/legacy/questions.py | 14 +- generalresearch/models/lucid/question.py | 12 +- generalresearch/models/lucid/survey.py | 6 +- generalresearch/models/morning/question.py | 6 +- generalresearch/models/morning/survey.py | 19 +- generalresearch/models/morning/task_collection.py | 6 +- generalresearch/models/network/__init__.py | 0 generalresearch/models/network/definitions.py | 70 ---- generalresearch/models/network/label.py | 129 ------ generalresearch/models/network/mtr/__init__.py | 0 generalresearch/models/network/mtr/command.py | 76 ---- generalresearch/models/network/mtr/execute.py | 60 --- generalresearch/models/network/mtr/parser.py | 19 - generalresearch/models/network/mtr/result.py | 175 --------- generalresearch/models/network/nmap/__init__.py | 0 generalresearch/models/network/nmap/command.py | 51 --- generalresearch/models/network/nmap/execute.py | 56 --- generalresearch/models/network/nmap/parser.py | 414 ------------------- generalresearch/models/network/nmap/result.py | 436 --------------------- generalresearch/models/network/rdns/__init__.py | 0 generalresearch/models/network/rdns/command.py | 38 -- generalresearch/models/network/rdns/execute.py | 44 --- generalresearch/models/network/rdns/parser.py | 21 - generalresearch/models/network/rdns/result.py | 52 --- generalresearch/models/network/tool_run.py | 121 ------ generalresearch/models/network/tool_run_command.py | 68 ---- generalresearch/models/network/utils.py | 5 - generalresearch/models/pollfish/question.py | 10 +- generalresearch/models/precision/question.py | 12 +- generalresearch/models/precision/survey.py | 7 +- .../models/precision/task_collection.py | 6 +- generalresearch/models/prodege/question.py | 14 +- generalresearch/models/prodege/survey.py | 11 +- generalresearch/models/prodege/task_collection.py | 6 +- generalresearch/models/repdata/question.py | 10 +- generalresearch/models/repdata/task_collection.py | 6 +- generalresearch/models/sago/question.py | 12 +- generalresearch/models/sago/survey.py | 20 +- generalresearch/models/sago/task_collection.py | 6 +- generalresearch/models/spectrum/question.py | 14 +- generalresearch/models/spectrum/survey.py | 4 +- generalresearch/models/spectrum/task_collection.py | 6 +- generalresearch/models/thl/contest/__init__.py | 14 +- generalresearch/models/thl/contest/contest.py | 10 +- .../models/thl/contest/contest_entry.py | 6 +- generalresearch/models/thl/contest/leaderboard.py | 40 +- generalresearch/models/thl/contest/raffle.py | 20 +- .../models/thl/profiling/upk_property.py | 5 +- generalresearch/models/thl/profiling/user_info.py | 14 +- generalresearch/models/thl/survey/__init__.py | 21 +- test_utils/managers/network/__init__.py | 0 test_utils/managers/network/conftest.py | 0 test_utils/models/network/__init__.py | 0 test_utils/models/network/conftest.py | 145 ------- tests/managers/network/__init__.py | 0 tests/managers/network/test_label.py | 209 ---------- tests/managers/network/test_tool_run.py | 25 -- tests/models/network/__init__.py | 0 tests/models/network/test_mtr.py | 33 -- tests/models/network/test_nmap.py | 39 -- tests/models/network/test_nmap_parser.py | 32 -- tests/models/network/test_rdns.py | 40 -- 77 files changed, 169 insertions(+), 3078 deletions(-) delete mode 100644 generalresearch/managers/network/__init__.py delete mode 100644 generalresearch/managers/network/label.py delete mode 100644 generalresearch/managers/network/mtr.py delete mode 100644 generalresearch/managers/network/nmap.py delete mode 100644 generalresearch/managers/network/rdns.py delete mode 100644 generalresearch/managers/network/tool_run.py delete mode 100644 generalresearch/models/network/__init__.py delete mode 100644 generalresearch/models/network/definitions.py delete mode 100644 generalresearch/models/network/label.py delete mode 100644 generalresearch/models/network/mtr/__init__.py delete mode 100644 generalresearch/models/network/mtr/command.py delete mode 100644 generalresearch/models/network/mtr/execute.py delete mode 100644 generalresearch/models/network/mtr/parser.py delete mode 100644 generalresearch/models/network/mtr/result.py delete mode 100644 generalresearch/models/network/nmap/__init__.py delete mode 100644 generalresearch/models/network/nmap/command.py delete mode 100644 generalresearch/models/network/nmap/execute.py delete mode 100644 generalresearch/models/network/nmap/parser.py delete mode 100644 generalresearch/models/network/nmap/result.py delete mode 100644 generalresearch/models/network/rdns/__init__.py delete mode 100644 generalresearch/models/network/rdns/command.py delete mode 100644 generalresearch/models/network/rdns/execute.py delete mode 100644 generalresearch/models/network/rdns/parser.py delete mode 100644 generalresearch/models/network/rdns/result.py delete mode 100644 generalresearch/models/network/tool_run.py delete mode 100644 generalresearch/models/network/tool_run_command.py delete mode 100644 generalresearch/models/network/utils.py delete mode 100644 test_utils/managers/network/__init__.py delete mode 100644 test_utils/managers/network/conftest.py delete mode 100644 test_utils/models/network/__init__.py delete mode 100644 test_utils/models/network/conftest.py delete mode 100644 tests/managers/network/__init__.py delete mode 100644 tests/managers/network/test_label.py delete mode 100644 tests/managers/network/test_tool_run.py delete mode 100644 tests/models/network/__init__.py delete mode 100644 tests/models/network/test_mtr.py delete mode 100644 tests/models/network/test_nmap.py delete mode 100644 tests/models/network/test_nmap_parser.py delete mode 100644 tests/models/network/test_rdns.py (limited to 'test_utils/models/network') diff --git a/generalresearch/managers/network/__init__.py b/generalresearch/managers/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/generalresearch/managers/network/label.py b/generalresearch/managers/network/label.py deleted file mode 100644 index cdea016..0000000 --- a/generalresearch/managers/network/label.py +++ /dev/null @@ -1,152 +0,0 @@ -from __future__ import annotations - -from collections.abc import Collection -from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING - -from psycopg import sql -from pydantic import TypeAdapter - -from generalresearch.managers.base import PostgresManager -from generalresearch.models.custom_types import ( - AwareDatetimeISO, - IPvAnyAddressStr, - IPvAnyNetwork, - IPvAnyNetworkStr, -) -from generalresearch.models.network.label import IPLabel - -if TYPE_CHECKING: - - from generalresearch.models.network.label import IPLabelKind, IPLabelSource - - -class IPLabelManager(PostgresManager): - def create(self, ip_label: IPLabel) -> IPLabel: - query = sql.SQL(""" - INSERT INTO network_iplabel ( - ip, labeled_at, created_at, - label_kind, source, confidence, - provider, metadata - ) VALUES ( - %(ip)s, %(labeled_at)s, %(created_at)s, - %(label_kind)s, %(source)s, %(confidence)s, - %(provider)s, %(metadata)s - ) RETURNING id;""") - params = ip_label.model_dump_postgres() - with self.pg_config.make_connection() as conn, conn.cursor() as c: - c.execute(query, params) - _pk = c.fetchone()["id"] - return ip_label - - def make_filter_str( - self, - ips: Collection[IPvAnyNetworkStr] | None = None, - ip_in_network: IPvAnyAddressStr | None = None, - label_kind: IPLabelKind | None = None, - source: IPLabelSource | None = None, - labeled_at: AwareDatetimeISO | None = None, - labeled_after: AwareDatetimeISO | None = None, - labeled_before: AwareDatetimeISO | None = None, - provider: str | None = None, - ): - filters = [] - params = {} - if labeled_after or labeled_before: - time_end = labeled_before or datetime.now(tz=UTC) - time_start = labeled_after or datetime(2017, 1, 1, tzinfo=UTC) - assert time_start.tzinfo.utcoffset(time_start) == timedelta(), "must be UTC" - assert time_end.tzinfo.utcoffset(time_end) == timedelta(), "must be UTC" - filters.append("labeled_at BETWEEN %(time_start)s AND %(time_end)s") - params["time_start"] = time_start - params["time_end"] = time_end - if labeled_at: - assert labeled_at.tzinfo.utcoffset(labeled_at) == timedelta(), "must be UTC" - filters.append("labeled_at == %(labeled_at)s") - params["labeled_at"] = labeled_at - if label_kind: - filters.append("label_kind = %(label_kind)s") - params["label_kind"] = label_kind.value - if source: - filters.append("source = %(source)s") - params["source"] = source.value - if provider: - filters.append("provider = %(provider)s") - params["provider"] = provider - if ips is not None: - filters.append("ip = ANY(%(ips)s)") - params["ips"] = list(ips) - if ip_in_network: - """ - Return matching networks. - e.g. ip = '13f9:c462:e039:a38c::1', might return rows - where ip = '13f9:c462:e039::/48' or '13f9:c462:e039:a38c::/64' - """ - filters.append("ip >>= %(ip_in_network)s") - params["ip_in_network"] = ip_in_network - - filter_str = "WHERE " + " AND ".join(filters) if filters else "" - return filter_str, params - - def filter( - self, - ips: Collection[IPvAnyNetworkStr] | None = None, - ip_in_network: IPvAnyAddressStr | None = None, - label_kind: IPLabelKind | None = None, - source: IPLabelSource | None = None, - labeled_at: AwareDatetimeISO | None = None, - labeled_after: AwareDatetimeISO | None = None, - labeled_before: AwareDatetimeISO | None = None, - provider: str | None = None, - ) -> list[IPLabel]: - filter_str, params = self.make_filter_str( - ips=ips, - ip_in_network=ip_in_network, - label_kind=label_kind, - source=source, - labeled_at=labeled_at, - labeled_after=labeled_after, - labeled_before=labeled_before, - provider=provider, - ) - query = f""" - SELECT - ip, labeled_at, created_at, - label_kind, source, confidence, - provider, metadata - FROM network_iplabel - {filter_str} - """ - res = self.pg_config.execute_sql_query(query, params) - return [IPLabel.model_validate(rec) for rec in res] - - def get_most_specific_matching_network(self, ip: IPvAnyAddressStr) -> IPvAnyNetwork: - """ - e.g. ip = 'b5f4:dc2:f136:70d5:5b6e:9a85:c7d4:3517', might return - 'b5f4:dc2:f136:70d5::/64' - """ - ip = TypeAdapter(IPvAnyAddressStr).validate_python(ip) - - query = """ - SELECT ip - FROM network_iplabel - WHERE ip >>= %(ip)s - ORDER BY masklen(ip) DESC - LIMIT 1;""" - res = self.pg_config.execute_sql_query(query, {"ip": ip}) - if res: - return IPvAnyNetwork(res[0]["ip"]) - - def test_join(self, ip): - query = """ - SELECT - to_jsonb(i) AS ipinfo, - to_jsonb(l) AS iplabel - FROM thl_ipinformation i - LEFT JOIN network_iplabel l - ON l.ip >>= i.ip - WHERE i.ip = %(ip)s - ORDER BY masklen(l.ip) DESC;""" - params = {"ip": ip} - res = self.pg_config.execute_sql_query(query, params) - return res diff --git a/generalresearch/managers/network/mtr.py b/generalresearch/managers/network/mtr.py deleted file mode 100644 index 7b79d96..0000000 --- a/generalresearch/managers/network/mtr.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from psycopg import Cursor, sql - -from generalresearch.managers.base import PostgresManager - -if TYPE_CHECKING: - from generalresearch.models.network.tool_run import MTRRun - - -class MTRRunManager(PostgresManager): - - def _create(self, run: MTRRun, c: Cursor | None = None) -> None: - """ - Do not use this directly. Must only be used in the context of a toolrun - """ - query = sql.SQL(""" - INSERT INTO network_mtr ( - run_id, source_ip, facility_id, - protocol, port, parsed, - started_at, ip, scan_group_id - ) - VALUES ( - %(run_id)s, %(source_ip)s, %(facility_id)s, - %(protocol)s, %(port)s, %(parsed)s, - %(started_at)s, %(ip)s, %(scan_group_id)s - ); - """) - params = run.model_dump_postgres() - - query_hops = sql.SQL(""" - INSERT INTO network_mtrhop ( - hop, ip, domain, asn, mtr_run_id - ) VALUES ( - %(hop)s, %(ip)s, %(domain)s, - %(asn)s, %(mtr_run_id)s - ) - """) - mtr_run = run.parsed - params_hops = [h.model_dump_postgres(run_id=run.id) for h in mtr_run.hops] - - if c: - c.execute(query, params) - if params_hops: - c.executemany(query_hops, params_hops) - - else: - with self.pg_config.make_connection() as conn, conn.cursor() as _c: - _c.execute(query, params) - if params_hops: - _c.executemany(query_hops, params_hops) diff --git a/generalresearch/managers/network/nmap.py b/generalresearch/managers/network/nmap.py deleted file mode 100644 index 96c6009..0000000 --- a/generalresearch/managers/network/nmap.py +++ /dev/null @@ -1,60 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from psycopg import Cursor, sql - -from generalresearch.managers.base import PostgresManager - -if TYPE_CHECKING: - from generalresearch.models.network.tool_run import NmapRun - - -class NmapRunManager(PostgresManager): - - def _create(self, run: NmapRun, c: Cursor | None = None) -> None: - """ - Insert a PortScan + PortScanPorts from a Pydantic NmapResult. - Do not use this directly. Must only be used in the context of a toolrun - """ - query = sql.SQL(""" - INSERT INTO network_portscan ( - run_id, xml_version, host_state, - host_state_reason, latency_ms, distance, - uptime_seconds, last_boot, - parsed, scan_group_id, open_tcp_ports, - started_at, ip, open_udp_ports - ) - VALUES ( - %(run_id)s, %(xml_version)s, %(host_state)s, - %(host_state_reason)s, %(latency_ms)s, %(distance)s, - %(uptime_seconds)s, %(last_boot)s, - %(parsed)s, %(scan_group_id)s, %(open_tcp_ports)s, - %(started_at)s, %(ip)s, %(open_udp_ports)s - ); - """) - params = run.model_dump_postgres() - - query_ports = sql.SQL(""" - INSERT INTO network_portscanport ( - port_scan_id, protocol, port, - state, reason, reason_ttl, - service_name - ) VALUES ( - %(port_scan_id)s, %(protocol)s, %(port)s, - %(state)s, %(reason)s, %(reason_ttl)s, - %(service_name)s - ) - """) - nmap_run = run.parsed - params_ports = [p.model_dump_postgres(run_id=run.id) for p in nmap_run.ports] - - if c: - c.execute(query, params) - if nmap_run.ports: - c.executemany(query_ports, params_ports) - else: - with self.pg_config.make_connection() as conn, conn.cursor(): - c.execute(query, params) - if nmap_run.ports: - c.executemany(query_ports, params_ports) diff --git a/generalresearch/managers/network/rdns.py b/generalresearch/managers/network/rdns.py deleted file mode 100644 index 1800364..0000000 --- a/generalresearch/managers/network/rdns.py +++ /dev/null @@ -1,37 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from psycopg import Cursor - -from generalresearch.managers.base import PostgresManager - -if TYPE_CHECKING: - from generalresearch.models.network.tool_run import RDNSRun - - -class RDNSRunManager(PostgresManager): - - def _create(self, run: RDNSRun, c: Cursor | None = None) -> None: - """ - Do not use this directly. Must only be used in the context of a toolrun - """ - query = """ - INSERT INTO network_rdnsresult ( - run_id, primary_hostname, primary_domain, - hostname_count, hostnames, - ip, started_at, scan_group_id - ) - VALUES ( - %(run_id)s, %(primary_hostname)s, %(primary_domain)s, - %(hostname_count)s, %(hostnames)s, - %(ip)s, %(started_at)s, %(scan_group_id)s - ); - """ - params = run.model_dump_postgres() - if c: - c.execute(query, params) - - else: - with self.pg_config.make_connection() as conn, conn.cursor() as _c: - _c.execute(query, params) diff --git a/generalresearch/managers/network/tool_run.py b/generalresearch/managers/network/tool_run.py deleted file mode 100644 index ec06305..0000000 --- a/generalresearch/managers/network/tool_run.py +++ /dev/null @@ -1,138 +0,0 @@ -from __future__ import annotations - -from collections.abc import Collection -from typing import TYPE_CHECKING - -from psycopg import Cursor, sql - -from generalresearch.managers.base import PostgresManager -from generalresearch.managers.network.mtr import MTRRunManager -from generalresearch.managers.network.nmap import NmapRunManager -from generalresearch.managers.network.rdns import RDNSRunManager -from generalresearch.models.network.rdns.result import RDNSResult -from generalresearch.models.network.tool_run import ( - MTRRun, - NmapRun, - RDNSRun, - ToolRun, -) - -if TYPE_CHECKING: - from generalresearch.managers.base import Permission - from generalresearch.models.network.tool_run import ToolName - from generalresearch.pg_helper import PostgresConfig - - -class ToolRunManager(PostgresManager): - def __init__( - self, - pg_config: PostgresConfig, - permissions: Collection[Permission] | None = None, - ): - super().__init__(pg_config=pg_config, permissions=permissions) - self.nmap_manager = NmapRunManager(self.pg_config) - self.rdns_manager = RDNSRunManager(self.pg_config) - self.mtr_manager = MTRRunManager(self.pg_config) - - def _create_tool_run(self, run: NmapRun | RDNSRun | MTRRun, c: Cursor): - query = sql.SQL(""" - INSERT INTO network_toolrun ( - ip, scan_group_id, tool_class, - tool_name, tool_version, started_at, - finished_at, status, raw_command, - config - ) - VALUES ( - %(ip)s, %(scan_group_id)s, %(tool_class)s, - %(tool_name)s, %(tool_version)s, %(started_at)s, - %(finished_at)s, %(status)s, %(raw_command)s, - %(config)s - ) RETURNING id; - """) - params = run.model_dump_postgres() - c.execute(query, params) - run_id = c.fetchone()["id"] - run.id = run_id - - def create_tool_run(self, run: NmapRun | RDNSRun | MTRRun): - if type(run) is NmapRun: - return self.create_nmap_run(run) - elif type(run) is RDNSRun: - return self.create_rdns_run(run) - elif type(run) is MTRRun: - return self.create_mtr_run(run) - else: - raise ValueError("unrecognized run type") - - def get_latest_runs_by_tool(self, ip: str) -> dict[ToolName, ToolRun]: - query = """ - SELECT DISTINCT ON (tool_name) * - FROM network_toolrun - WHERE ip = %(ip)s - ORDER BY tool_name, started_at DESC; - """ - params = {"ip": ip} - res = self.pg_config.execute_sql_query(query, params=params) - runs = [ToolRun.model_validate(x) for x in res] - return {r.tool_name: r for r in runs} - - def create_nmap_run(self, run: NmapRun) -> NmapRun: - """ - Insert a PortScan + PortScanPorts from a Pydantic NmapResult. - """ - with self.pg_config.make_connection() as conn, conn.cursor() as c: - self._create_tool_run(run, c) - self.nmap_manager._create(run, c=c) - return run - - def get_nmap_run(self, id: int) -> NmapRun: - query = """ - SELECT tr.*, np.parsed - FROM network_toolrun tr - JOIN network_portscan np ON tr.id = np.run_id - WHERE id = %(id)s - """ - params = {"id": id} - res = self.pg_config.execute_sql_query(query, params)[0] - return NmapRun.model_validate(res) - - def create_rdns_run(self, run: RDNSRun) -> RDNSRun: - """ - Insert a RDnsRun + RDNSResult - """ - with self.pg_config.make_connection() as conn, conn.cursor() as c: - self._create_tool_run(run, c) - self.rdns_manager._create(run, c=c) - return run - - def get_rdns_run(self, id: int) -> RDNSRun: - query = """ - SELECT tr.*, hostnames - FROM network_toolrun tr - JOIN network_rdnsresult np ON tr.id = np.run_id - WHERE id = %(id)s - """ - params = {"id": id} - res = self.pg_config.execute_sql_query(query, params)[0] - parsed = RDNSResult.model_validate( - {"ip": res["ip"], "hostnames": res["hostnames"]} - ) - res["parsed"] = parsed - return RDNSRun.model_validate(res) - - def create_mtr_run(self, run: MTRRun) -> MTRRun: - with self.pg_config.make_connection() as conn, conn.cursor() as c: - self._create_tool_run(run, c) - self.mtr_manager._create(run, c=c) - return run - - def get_mtr_run(self, id: int) -> MTRRun: - query = """ - SELECT tr.*, mtr.parsed, mtr.source_ip, mtr.facility_id - FROM network_toolrun tr - JOIN network_mtr mtr ON tr.id = mtr.run_id - WHERE id = %(id)s - """ - params = {"id": id} - res = self.pg_config.execute_sql_query(query, params)[0] - return MTRRun.model_validate(res) diff --git a/generalresearch/models/cint/question.py b/generalresearch/models/cint/question.py index 0f5453b..89c7871 100644 --- a/generalresearch/models/cint/question.py +++ b/generalresearch/models/cint/question.py @@ -3,11 +3,12 @@ from __future__ import annotations import json from datetime import UTC, datetime from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator +from generalresearch.models.cint import CintQuestionIdType from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.definitions import Source from generalresearch.models.string_utils import remove_nbsp @@ -15,12 +16,9 @@ from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, ) - -if TYPE_CHECKING: - from generalresearch.models.cint import CintQuestionIdType - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) class CintQuestionType(StrEnum): diff --git a/generalresearch/models/cint/survey.py b/generalresearch/models/cint/survey.py index cd429dd..b2a8935 100644 --- a/generalresearch/models/cint/survey.py +++ b/generalresearch/models/cint/survey.py @@ -4,7 +4,7 @@ import json import logging from datetime import UTC, datetime from decimal import Decimal -from typing import TYPE_CHECKING, Annotated, Any, Literal, Self +from typing import Annotated, Any, Literal, Self from more_itertools import flatten from pydantic import ( @@ -18,6 +18,7 @@ from pydantic import ( ) from generalresearch.locales import Localelator +from generalresearch.models.cint import CintQuestionIdType from generalresearch.models.custom_types import ( AlphaNumStr, AwareDatetimeISO, @@ -31,10 +32,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - from generalresearch.models.cint import CintQuestionIdType - - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -76,9 +73,9 @@ class CintQuota(BaseModel): @model_validator(mode="after") def validate_condition_len(self) -> Self: if self.quota_type == "total": - assert ( - self.condition_hashes is None - ), "total quota should not have conditions" + assert self.condition_hashes is None, ( + "total quota should not have conditions" + ) elif self.quota_type == "client": assert len(self.condition_hashes) > 0, "quota must have conditions" return self diff --git a/generalresearch/models/cint/task_collection.py b/generalresearch/models/cint/task_collection.py index 31a0173..4ae8de4 100644 --- a/generalresearch/models/cint/task_collection.py +++ b/generalresearch/models/cint/task_collection.py @@ -1,19 +1,15 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator +from generalresearch.models.cint.survey import CintSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.cint.survey import CintSurvey - COUNTRY_ISOS: set[str] = Localelator().get_all_countries() LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() diff --git a/generalresearch/models/dynata/survey.py b/generalresearch/models/dynata/survey.py index 491157b..6ff397f 100644 --- a/generalresearch/models/dynata/survey.py +++ b/generalresearch/models/dynata/survey.py @@ -5,7 +5,7 @@ import logging from datetime import UTC from decimal import Decimal from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from more_itertools import flatten from pydantic import ( @@ -26,7 +26,7 @@ from generalresearch.models.custom_types import ( CoercedStr, DeviceTypes, ) -from generalresearch.models.definitions import Source +from generalresearch.models.definitions import Source, TaskCalculationType from generalresearch.models.dynata import DynataStatus from generalresearch.models.thl.demographics import ( Gender, @@ -37,10 +37,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - - from generalresearch.models.definitions import TaskCalculationType - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -136,7 +132,7 @@ class DynataCondition(MarketplaceCondition): if cell["kind"] == "RANGE": d["values"] = [ - f"{cell["range"]["from"] or "inf"}-{cell["range"]["to"] or "inf"}" + f"{cell['range']['from'] or 'inf'}-{cell['range']['to'] or 'inf'}" ] d["value_type"] = ConditionValueType.RANGE return cls.model_validate(d) diff --git a/generalresearch/models/dynata/task_collection.py b/generalresearch/models/dynata/task_collection.py index c6cdc19..e6f0548 100644 --- a/generalresearch/models/dynata/task_collection.py +++ b/generalresearch/models/dynata/task_collection.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index @@ -8,14 +8,12 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.dynata import DynataStatus +from generalresearch.models.dynata.survey import DynataSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.dynata.survey import DynataSurvey - COUNTRY_ISOS = Localelator().get_all_countries() LANGUAGE_ISOS = Localelator().get_all_languages() diff --git a/generalresearch/models/innovate/question.py b/generalresearch/models/innovate/question.py index 6423399..4af0639 100644 --- a/generalresearch/models/innovate/question.py +++ b/generalresearch/models/innovate/question.py @@ -4,21 +4,19 @@ from __future__ import annotations import json import logging from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator from generalresearch.models.definitions import Source +from generalresearch.models.innovate import InnovateQuestionID from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, ) - -if TYPE_CHECKING: - from generalresearch.models.innovate import InnovateQuestionID - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/innovate/survey.py b/generalresearch/models/innovate/survey.py index 7232228..3ca990c 100644 --- a/generalresearch/models/innovate/survey.py +++ b/generalresearch/models/innovate/survey.py @@ -6,7 +6,6 @@ from datetime import UTC, date from decimal import Decimal from functools import cached_property from typing import ( - TYPE_CHECKING, Annotated, Any, Literal, @@ -33,12 +32,14 @@ from generalresearch.models.custom_types import ( from generalresearch.models.definitions import ( LogicalOperator, Source, + TaskCalculationType, ) from generalresearch.models.innovate import ( InnovateDuplicateCheckLevel, InnovateQuotaStatus, InnovateStatus, ) +from generalresearch.models.innovate.question import InnovateQuestionID from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask from generalresearch.models.thl.survey.condition import ( @@ -46,13 +47,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - - from generalresearch.models.definitions import ( - TaskCalculationType, - ) - from generalresearch.models.innovate.question import InnovateQuestionID - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) diff --git a/generalresearch/models/innovate/task_collection.py b/generalresearch/models/innovate/task_collection.py index a647d30..7bf9d0f 100644 --- a/generalresearch/models/innovate/task_collection.py +++ b/generalresearch/models/innovate/task_collection.py @@ -1,20 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.innovate import InnovateStatus +from generalresearch.models.innovate.survey import InnovateSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.innovate.survey import InnovateSurvey - COUNTRY_ISOS: set[str] = Localelator().get_all_countries() LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() diff --git a/generalresearch/models/legacy/offerwall.py b/generalresearch/models/legacy/offerwall.py index 013b506..da28663 100644 --- a/generalresearch/models/legacy/offerwall.py +++ b/generalresearch/models/legacy/offerwall.py @@ -1,34 +1,27 @@ from __future__ import annotations -from typing import TYPE_CHECKING - from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.legacy.bucket import ( + BucketBase, + MarketplaceBucket, + OneShotOfferwallBucket, + OneShotSoftPairOfferwallBucket, + SingleEntryBucket, + SoftPairBucket, + TimeBucksBucket, + TopNBucket, + TopNPlusBucket, + TopNPlusRecontactBucket, + WXETOfferwallBucket, +) from generalresearch.models.legacy.definitions import OfferwallReason from generalresearch.models.thl.payout_format import ( PayoutFormatField, + PayoutFormatType, ) - -if TYPE_CHECKING: - - from generalresearch.models.legacy.bucket import ( - BucketBase, - MarketplaceBucket, - OneShotOfferwallBucket, - OneShotSoftPairOfferwallBucket, - SingleEntryBucket, - SoftPairBucket, - TimeBucksBucket, - TopNBucket, - TopNPlusBucket, - TopNPlusRecontactBucket, - WXETOfferwallBucket, - ) - from generalresearch.models.thl.payout_format import ( - PayoutFormatType, - ) - from generalresearch.models.thl.profiling.upk_question import UpkQuestion +from generalresearch.models.thl.profiling.upk_question import UpkQuestion """ Not Done: diff --git a/generalresearch/models/legacy/questions.py b/generalresearch/models/legacy/questions.py index e6803f0..caa6aae 100644 --- a/generalresearch/models/legacy/questions.py +++ b/generalresearch/models/legacy/questions.py @@ -17,17 +17,15 @@ from sentry_sdk import capture_exception from generalresearch.models.custom_types import UUIDStr from generalresearch.models.legacy.api_status import StatusResponse +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestionOut, +) +from generalresearch.models.thl.session import Wall +from generalresearch.models.thl.user import User if TYPE_CHECKING: - from generalresearch.managers.thl.user_manager.user_manager import ( - UserManager, - ) + from generalresearch.managers.thl.user_manager.user_manager import UserManager from generalresearch.managers.thl.wall import WallManager - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestionOut, - ) - from generalresearch.models.thl.session import Wall - from generalresearch.models.thl.user import User class UpkQuestionResponse(StatusResponse): diff --git a/generalresearch/models/lucid/question.py b/generalresearch/models/lucid/question.py index c1b9e52..288f0d2 100644 --- a/generalresearch/models/lucid/question.py +++ b/generalresearch/models/lucid/question.py @@ -2,20 +2,18 @@ from __future__ import annotations import logging from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from pydantic import BaseModel, Field, field_validator, model_validator from generalresearch.models.definitions import Source +from generalresearch.models.lucid import LucidQuestionIdType from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, ) - -if TYPE_CHECKING: - from generalresearch.models.lucid import LucidQuestionIdType - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/lucid/survey.py b/generalresearch/models/lucid/survey.py index 02b31ab..bca471b 100644 --- a/generalresearch/models/lucid/survey.py +++ b/generalresearch/models/lucid/survey.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Self +from typing import Any, Self from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt @@ -11,14 +11,12 @@ from generalresearch.models.custom_types import ( UUIDStr, ) from generalresearch.models.definitions import Source +from generalresearch.models.thl.locales import CountryISO, LanguageISO from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, ) -if TYPE_CHECKING: - from generalresearch.models.thl.locales import CountryISO, LanguageISO - class LucidCondition(MarketplaceCondition): model_config = ConfigDict(populate_by_name=True, frozen=False, extra="ignore") diff --git a/generalresearch/models/morning/question.py b/generalresearch/models/morning/question.py index 909992f..7c676fb 100644 --- a/generalresearch/models/morning/question.py +++ b/generalresearch/models/morning/question.py @@ -1,20 +1,18 @@ import json from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator from generalresearch.locales import Localelator from generalresearch.models.definitions import Source +from generalresearch.models.morning import MorningQuestionID from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, ) -if TYPE_CHECKING: - from generalresearch.models.morning import MorningQuestionID - # todo: we could validate that the country_iso / language_iso exists ... locale_helper = Localelator() diff --git a/generalresearch/models/morning/survey.py b/generalresearch/models/morning/survey.py index 01cc4ff..8738631 100644 --- a/generalresearch/models/morning/survey.py +++ b/generalresearch/models/morning/survey.py @@ -6,7 +6,6 @@ from datetime import UTC from decimal import Decimal from functools import cached_property from typing import ( - TYPE_CHECKING, Annotated, Any, Literal, @@ -30,24 +29,20 @@ from generalresearch.models.custom_types import ( UUIDStrCoerce, ) from generalresearch.models.definitions import Source -from generalresearch.models.morning import MorningStatus +from generalresearch.models.morning import MorningQuestionID, MorningStatus +from generalresearch.models.morning.question import MorningQuestion from generalresearch.models.thl.demographics import Gender +from generalresearch.models.thl.locales import ( + CountryISO, + CountryISOs, + LanguageISOs, +) from generalresearch.models.thl.survey import MarketplaceTask from generalresearch.models.thl.survey.condition import ( ConditionValueType, MarketplaceCondition, ) -if TYPE_CHECKING: - - from generalresearch.models.morning import MorningQuestionID - from generalresearch.models.morning.question import MorningQuestion - from generalresearch.models.thl.locales import ( - CountryISO, - CountryISOs, - LanguageISOs, - ) - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) diff --git a/generalresearch/models/morning/task_collection.py b/generalresearch/models/morning/task_collection.py index 1117937..9303a2f 100644 --- a/generalresearch/models/morning/task_collection.py +++ b/generalresearch/models/morning/task_collection.py @@ -1,20 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.morning import MorningStatus +from generalresearch.models.morning.survey import MorningBid from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.morning.survey import MorningBid - COUNTRY_ISOS: set[str] = Localelator().get_all_countries() LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() diff --git a/generalresearch/models/network/__init__.py b/generalresearch/models/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/generalresearch/models/network/definitions.py b/generalresearch/models/network/definitions.py deleted file mode 100644 index 2e1ab91..0000000 --- a/generalresearch/models/network/definitions.py +++ /dev/null @@ -1,70 +0,0 @@ -from __future__ import annotations - -from enum import StrEnum -from ipaddress import ip_address, ip_network - -CGNAT_NET = ip_network("100.64.0.0/10") - - -class IPProtocol(StrEnum): - TCP = "tcp" - UDP = "udp" - SCTP = "sctp" - IP = "ip" - ICMP = "icmp" - ICMPv6 = "icmpv6" - - def to_number(self) -> int: - # https://www.iana.org/assignments/protocol-numbers/protocol-numbers.xhtml - return { - self.TCP: 6, - self.UDP: 17, - self.SCTP: 132, - self.IP: 4, - self.ICMP: 1, - self.ICMPv6: 58, - }[self] - - -class IPKind(StrEnum): - PUBLIC = "public" - PRIVATE = "private" - CGNAT = "carrier_nat" - LOOPBACK = "loopback" - LINK_LOCAL = "link_local" - MULTICAST = "multicast" - RESERVED = "reserved" - UNSPECIFIED = "unspecified" - - -def get_ip_kind(ip: str | None) -> IPKind | None: - if not ip: - return None - - ip_obj = ip_address(ip) - - if ip_obj in CGNAT_NET: - return IPKind.CGNAT - - if ip_obj.is_loopback: - return IPKind.LOOPBACK - - if ip_obj.is_link_local: - return IPKind.LINK_LOCAL - - if ip_obj.is_multicast: - return IPKind.MULTICAST - - if ip_obj.is_unspecified: - return IPKind.UNSPECIFIED - - if ip_obj.is_private: - return IPKind.PRIVATE - - if ip_obj.is_reserved: - return IPKind.RESERVED - - if ip_obj.is_global: - return IPKind.PUBLIC - - return None diff --git a/generalresearch/models/network/label.py b/generalresearch/models/network/label.py deleted file mode 100644 index c27f36f..0000000 --- a/generalresearch/models/network/label.py +++ /dev/null @@ -1,129 +0,0 @@ -from __future__ import annotations - -import ipaddress -from enum import StrEnum -from ipaddress import IPv4Network, IPv6Network -from typing import TYPE_CHECKING - -from pydantic import ( - BaseModel, - ConfigDict, - Field, - IPvAnyNetwork, - computed_field, - field_validator, -) - -from generalresearch.models.custom_types import now_utc_factory - -if TYPE_CHECKING: - from generalresearch.models.custom_types import AwareDatetimeISO - - -class IPTrustClass(StrEnum): - TRUSTED = "trusted" - UNTRUSTED = "untrusted" - # Note: use case of unknown is for e.g. Spur says this IP is a residential proxy - # on 2026-1-1, and then has no annotation a month later. It doesn't mean - # the IP is TRUSTED, but we want to record that Spur now doesn't claim UNTRUSTED. - UNKNOWN = "unknown" - - -class IPLabelKind(StrEnum): - # --- UNTRUSTED --- - RESIDENTIAL_PROXY = "residential_proxy" - DATACENTER_PROXY = "datacenter_proxy" - ISP_PROXY = "isp_proxy" - MOBILE_PROXY = "mobile_proxy" - PROXY = "proxy" - HOSTING = "hosting" - VPN = "vpn" - RELAY = "relay" - TOR_EXIT = "tor_exit" - BAD_ACTOR = "bad_actor" - # --- TRUSTED --- - TRUSTED_USER = "trusted_user" - # --- UNKNOWN --- - UNKNOWN = "unknown" - - -class IPLabelSource(StrEnum): - # We got this IP from our own use of a proxy service - INTERNAL_USE = "internal_use" - - # An external "security" service flagged this IP - GRIP = "grip" - SPUR = "spur" - IPINFO = "ipinfo" - MAXMIND = "maxmind" - - MANUAL = "manual" - - -class IPLabel(BaseModel): - """ - Stores *ground truth* about an IP at a specific time. - To be used for model training and evaluation. - """ - - model_config = ConfigDict(validate_assignment=True) - - ip: IPvAnyNetwork = Field() - - labeled_at: AwareDatetimeISO = Field(default_factory=now_utc_factory) - created_at: AwareDatetimeISO | None = Field(default=None) - - label_kind: IPLabelKind = Field() - source: IPLabelSource = Field() - - confidence: float = Field(default=1.0, ge=0.0, le=1.0) - - # Optionally, if this is untrusted, which service is providing the proxy/vpn service - provider: str | None = Field( - default=None, examples=["geonode", "gecko"], max_length=128 - ) - - metadata: IPLabelMetadata | None = Field(default=None) - - @field_validator("ip", mode="before") - @classmethod - def normalize_and_validate_network( - cls, v: IPvAnyNetwork - ) -> IPv4Network | IPv6Network | None: - net = ipaddress.ip_network(address=v, strict=False) - - if isinstance(net, ipaddress.IPv6Network) and net.prefixlen > 64: - raise ValueError("IPv6 network must be /64 or larger") - - return net - - @field_validator("provider", mode="before") - @classmethod - def provider_format(cls, v: str | None) -> str | None: - if v is None: - return v - return v.lower().strip() - - @computed_field() - @property - def trust_class(self) -> IPTrustClass: - if self.label_kind == IPLabelKind.UNKNOWN: - return IPTrustClass.UNKNOWN - if self.label_kind == IPLabelKind.TRUSTED_USER: - return IPTrustClass.TRUSTED - return IPTrustClass.UNTRUSTED - - def model_dump_postgres(self): - d = self.model_dump(mode="json") - d["metadata"] = self.metadata.model_dump_json() if self.metadata else None - return d - - -class IPLabelMetadata(BaseModel): - """ - To be expanded. Just for storing some things from Spur for now - """ - - model_config = ConfigDict(validate_assignment=True, extra="allow") - - services: list[str] | None = Field(min_length=1, examples=[["RDP"]]) diff --git a/generalresearch/models/network/mtr/__init__.py b/generalresearch/models/network/mtr/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/generalresearch/models/network/mtr/command.py b/generalresearch/models/network/mtr/command.py deleted file mode 100644 index 7e74f20..0000000 --- a/generalresearch/models/network/mtr/command.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -import subprocess -from typing import TYPE_CHECKING - -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.mtr.parser import parse_mtr_output - -if TYPE_CHECKING: - from generalresearch.models.network.mtr.result import MTRResult - from generalresearch.models.network.tool_run_command import MTRRunCommand - -SUPPORTED_PROTOCOLS = { - IPProtocol.TCP, - IPProtocol.UDP, - IPProtocol.SCTP, - IPProtocol.ICMP, -} -PROTOCOLS_W_PORT = {IPProtocol.TCP, IPProtocol.UDP, IPProtocol.SCTP} - - -def build_mtr_command( - ip: str, - protocol: IPProtocol | None = None, - port: int | None = None, - report_cycles: int | None = 10, -) -> str: - # https://manpages.ubuntu.com/manpages/focal/man8/mtr.8.html - # e.g. "mtr -r -c 2 -b -z -j -T -P 443 74.139.70.149" - args = ["mtr", "--report", "--show-ips", "--aslookup", "--json"] - if report_cycles is not None: - args.extend(["-c", str(int(report_cycles))]) - if port is not None: - if protocol is None: - protocol = IPProtocol.TCP - assert protocol in PROTOCOLS_W_PORT, "port only allowed for TCP/SCTP/UDP traces" - args.extend(["--port", str(int(port))]) - if protocol: - assert protocol in SUPPORTED_PROTOCOLS, f"unsupported protocol: {protocol}" - # default is ICMP (no args) - arg_map = { - IPProtocol.TCP: "--tcp", - IPProtocol.UDP: "--udp", - IPProtocol.SCTP: "--sctp", - } - if protocol in arg_map: - args.append(arg_map[protocol]) - args.append(ip) - return " ".join(args) - - -def get_mtr_version() -> str: - proc = subprocess.run( - ["mtr", "-v"], - capture_output=True, - text=True, - check=False, - ) - # e.g. mtr 0.95 - ver_str = proc.stdout.strip() - return ver_str.split(" ", 1)[1] - - -def run_mtr(config: MTRRunCommand) -> MTRResult: - cmd = config.to_command_str() - args = cmd.split(" ") - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - return parse_mtr_output( - raw, protocol=config.options.protocol, port=config.options.port - ) diff --git a/generalresearch/models/network/mtr/execute.py b/generalresearch/models/network/mtr/execute.py deleted file mode 100644 index 5ab7632..0000000 --- a/generalresearch/models/network/mtr/execute.py +++ /dev/null @@ -1,60 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from uuid import uuid4 - -from generalresearch.models.custom_types import UUIDStr -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.mtr.command import ( - get_mtr_version, - run_mtr, -) -from generalresearch.models.network.tool_run import ( - MTRRun, - Status, - ToolClass, - ToolName, -) -from generalresearch.models.network.tool_run_command import ( - MTRRunCommand, - MTRRunCommandOptions, -) -from generalresearch.models.network.utils import get_source_ip - - -def execute_mtr( - ip: str, - scan_group_id: UUIDStr | None = None, - protocol: IPProtocol | None = IPProtocol.ICMP, - port: int | None = None, - report_cycles: int = 10, -) -> MTRRun: - config = MTRRunCommand( - options=MTRRunCommandOptions( - ip=ip, - report_cycles=report_cycles, - protocol=protocol, - port=port, - ), - ) - - started_at = datetime.now(tz=UTC) - tool_version = get_mtr_version() - result = run_mtr(config) - finished_at = datetime.now(tz=UTC) - - return MTRRun( - tool_name=ToolName.MTR, - tool_class=ToolClass.TRACEROUTE, - tool_version=tool_version, - status=Status.SUCCESS, - ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=config.to_command_str(), - scan_group_id=scan_group_id or uuid4().hex, - config=config, - parsed=result, - source_ip=get_source_ip(), - facility_id=1, - ) diff --git a/generalresearch/models/network/mtr/parser.py b/generalresearch/models/network/mtr/parser.py deleted file mode 100644 index 30c22bf..0000000 --- a/generalresearch/models/network/mtr/parser.py +++ /dev/null @@ -1,19 +0,0 @@ -import json -from typing import Any - -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.mtr.result import MTRResult - - -def parse_mtr_output(raw: str, port: int, protocol: IPProtocol) -> MTRResult: - data = parse_mtr_raw_output(raw) - data["port"] = port - data["protocol"] = protocol - return MTRResult.model_validate(data) - - -def parse_mtr_raw_output(raw: str) -> dict[str, Any]: - data = json.loads(raw)["report"] - data.update(data.pop("mtr")) - data["hops"] = data.pop("hubs") - return data diff --git a/generalresearch/models/network/mtr/result.py b/generalresearch/models/network/mtr/result.py deleted file mode 100644 index d17136c..0000000 --- a/generalresearch/models/network/mtr/result.py +++ /dev/null @@ -1,175 +0,0 @@ -from __future__ import annotations - -import re -from functools import cached_property -from ipaddress import ip_address -from typing import TYPE_CHECKING - -import tldextract -from pydantic import ( - BaseModel, - ConfigDict, - Field, - computed_field, - field_validator, - model_validator, -) - -from generalresearch.models.network.definitions import ( - IPProtocol, - get_ip_kind, -) - -if TYPE_CHECKING: - from generalresearch.models.network.definitions import IPKind - -HOST_RE = re.compile(r"^(?P.+?) \((?P[^)]+)\)$") - - -class MTRHop(BaseModel): - model_config = ConfigDict(populate_by_name=True) - - hop: int = Field(alias="count") - host: str - asn: int | None = Field(default=None, alias="ASN") - - loss_pct: float = Field(alias="Loss%") - sent: int = Field(alias="Snt") - - last_ms: float = Field(alias="Last") - avg_ms: float = Field(alias="Avg") - best_ms: float = Field(alias="Best") - worst_ms: float = Field(alias="Wrst") - stdev_ms: float = Field(alias="StDev") - - hostname: str | None = Field( - default=None, examples=["fixed-187-191-8-145.totalplay.net"] - ) - ip: str | None = None - - @field_validator("asn", mode="before") - @classmethod - def normalize_asn(cls, v: str): - if v is None or v == "AS???": - return None - if type(v) is int: - return v - return int(v.replace("AS", "")) - - @model_validator(mode="after") - def parse_host(self): - host = self.host.strip() - - # hostname (ip) - m = HOST_RE.match(host) - if m: - self.hostname = m.group("hostname") - self.ip = m.group("ip") - return self - - # ip only - try: - ip_address(host) - self.ip = host - self.hostname = None - return self - except ValueError: - pass - - # hostname only - self.hostname = host - self.ip = None - return self - - @cached_property - def ip_kind(self) -> IPKind | None: - return get_ip_kind(self.ip) - - @cached_property - def icmp_rate_limited(self): - if self.avg_ms == 0: - return False - return self.stdev_ms > self.avg_ms or self.worst_ms > self.best_ms * 10 - - @computed_field(examples=["totalplay.net"]) - @cached_property - def domain(self) -> str | None: - if self.hostname: - return tldextract.extract(self.hostname).top_domain_under_public_suffix - - def model_dump_postgres(self, run_id: int): - # Writes for the network_mtrhop table - d = {"mtr_run_id": run_id} - data = self.model_dump( - mode="json", - include={ - "hop", - "ip", - "domain", - "asn", - }, - ) - d.update(data) - return d - - -class MTRResult(BaseModel): - model_config = ConfigDict(populate_by_name=True) - - source: str = Field(description="Hostname of the system running mtr.", alias="src") - destination: str = Field( - description="Destination hostname or IP being traced.", alias="dst" - ) - tos: int = Field(description="IP Type-of-Service (TOS) value used for probes.") - tests: int = Field(description="Number of probes sent per hop.") - psize: int = Field(description="Probe packet size in bytes.") - bitpattern: str = Field(description="Payload byte pattern used in probes (hex).") - - # Protocol used for the traceroute - protocol: IPProtocol = Field(default=IPProtocol.ICMP) - # The target port number for TCP/SCTP/UDP traces - port: int | None = Field(default=None) - - hops: list[MTRHop] = Field() - - def model_dump_postgres(self): - # Writes for the network_mtr table - d = self.model_dump( - mode="json", - include={"port"}, - ) - d["protocol"] = self.protocol.to_number() - d["parsed"] = self.model_dump_json(indent=0) - return d - - def print_report(self) -> None: - print( - f"MTR Report → {self.destination} {self.protocol.name} {self.port or ''}\n" - ) - host_max_len = max(len(h.host) for h in self.hops) - - header = ( - f"{'Hop':>3} " - f"{'Host':<{host_max_len}} " - f"{'Kind':<10} " - f"{'ASN':<8} " - f"{'Loss%':>6} {'Sent':>5} " - f"{'Last':>7} {'Avg':>7} {'Best':>7} {'Worst':>7} {'StDev':>7}" - ) - print(header) - print("-" * len(header)) - - for hop in self.hops: - print( - f"{hop.hop:>3} " - f"{hop.host:<{host_max_len}} " - f"{hop.ip_kind or '???':<10} " - f"{hop.asn or '???':<8} " - f"{hop.loss_pct:6.1f} " - f"{hop.sent:5d} " - f"{hop.last_ms:7.1f} " - f"{hop.avg_ms:7.1f} " - f"{hop.best_ms:7.1f} " - f"{hop.worst_ms:7.1f} " - f"{hop.stdev_ms:7.1f}" - ) diff --git a/generalresearch/models/network/nmap/__init__.py b/generalresearch/models/network/nmap/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/generalresearch/models/network/nmap/command.py b/generalresearch/models/network/nmap/command.py deleted file mode 100644 index 3509b8d..0000000 --- a/generalresearch/models/network/nmap/command.py +++ /dev/null @@ -1,51 +0,0 @@ -from __future__ import annotations - -import subprocess -from typing import TYPE_CHECKING - -from generalresearch.models.network.nmap.parser import parse_nmap_xml - -if TYPE_CHECKING: - from generalresearch.models.network.nmap.result import NmapResult - from generalresearch.models.network.tool_run_command import NmapRunCommand - - -def build_nmap_command( - ip: str, - no_ping: bool = True, - enable_advanced: bool = True, - timing: int = 4, - ports: str | None = None, - top_ports: int | None = None, -) -> str: - # e.g. "nmap -Pn -T4 -A --top-ports 1000 -oX - scanme.nmap.org" - # https://linux.die.net/man/1/nmap - args = ["nmap"] - assert 0 <= timing <= 5 - args.append(f"-T{timing}") - if no_ping: - args.append("-Pn") - if enable_advanced: - args.append("-A") - if ports is not None: - assert top_ports is None - args.extend(["-p", ports]) - if top_ports is not None: - assert ports is None - args.extend(["--top-ports", str(top_ports)]) - - args.extend(["-oX", "-", ip]) - return " ".join(args) - - -def run_nmap(config: NmapRunCommand) -> NmapResult: - cmd = config.to_command_str() - args = cmd.split(" ") - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - return parse_nmap_xml(raw) diff --git a/generalresearch/models/network/nmap/execute.py b/generalresearch/models/network/nmap/execute.py deleted file mode 100644 index 09ec28b..0000000 --- a/generalresearch/models/network/nmap/execute.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -from uuid import uuid4 - -from generalresearch.models.custom_types import UUIDStr -from generalresearch.models.network.nmap.command import run_nmap -from generalresearch.models.network.tool_run import ( - NmapRun, - Status, - ToolClass, - ToolName, -) -from generalresearch.models.network.tool_run_command import ( - NmapRunCommand, - NmapRunCommandOptions, -) - - -def execute_nmap( - ip: str, - top_ports: int | None = 1000, - ports: str | None = None, - no_ping: bool = True, - enable_advanced: bool = True, - timing: int = 4, - scan_group_id: UUIDStr | None = None, -) -> NmapRun: - config = NmapRunCommand( - options=NmapRunCommandOptions( - top_ports=top_ports, - ports=ports, - no_ping=no_ping, - enable_advanced=enable_advanced, - timing=timing, - ip=ip, - ) - ) - result = run_nmap(config) - assert result.exit_status == "success" - assert result.target_ip == ip, f"{result.target_ip=}, {ip=}" - assert result.command_line == config.to_command_str() - - run = NmapRun( - tool_name=ToolName.NMAP, - tool_class=ToolClass.PORT_SCAN, - tool_version=result.version, - status=Status.SUCCESS, - ip=ip, - started_at=result.started_at, - finished_at=result.finished_at, - raw_command=result.command_line, - scan_group_id=scan_group_id or uuid4().hex, - config=config, - parsed=result, - ) - return run diff --git a/generalresearch/models/network/nmap/parser.py b/generalresearch/models/network/nmap/parser.py deleted file mode 100644 index 866b4bd..0000000 --- a/generalresearch/models/network/nmap/parser.py +++ /dev/null @@ -1,414 +0,0 @@ -from __future__ import annotations - -import xml.etree.ElementTree as ET -from datetime import UTC, datetime -from typing import Any - -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.nmap.result import ( - NmapHostname, - NmapHostScript, - NmapHostState, - NmapHostStatusReason, - NmapOSClass, - NmapOSMatch, - NmapPort, - NmapPortStats, - NmapResult, - NmapScanInfo, - NmapScanType, - NmapScript, - NmapService, - NmapTrace, - NmapTraceHop, - PortState, - PortStateReason, -) - - -class NmapParserException(Exception): - def __init__(self, msg): - self.msg = msg - - def __str__(self): - return self.msg - - -class NmapXmlParser: - """ - Example: https://nmap.org/book/output-formats-xml-output.html - Full DTD: https://nmap.org/book/nmap-dtd.html - """ - - @classmethod - def parse_xml(cls, nmap_data: str) -> NmapResult: - """ - Expects a full nmap scan report. - """ - - try: - root = ET.fromstring(nmap_data) - except ET.ParseError as e: - emsg = f"Wrong XML structure: cannot parse data: {e}" - raise NmapParserException(emsg) - - if root.tag != "nmaprun": - raise NmapParserException("Unpexpected data structure for XML " "root node") - return cls._parse_xml_nmaprun(root) - - @classmethod - def _parse_xml_nmaprun(cls, root: ET.Element) -> NmapResult: - """ - This method parses out a full nmap scan report from its XML root - node: . We expect there is only 1 host in this report! - - :param root: Element from xml.ElementTree (top of XML the document) - """ - cls._validate_nmap_root(root) - host_count = len(root.findall(".//host")) - assert host_count == 1, f"Expected 1 host, got {host_count}" - - xml_str = ET.tostring(root, encoding="unicode").replace("\n", "") - nmap_data = {"raw_xml": xml_str} - nmap_data.update(cls._parse_nmaprun(root)) - - nmap_data["scan_infos"] = [ - cls._parse_scaninfo(scaninfo_el) - for scaninfo_el in root.findall(".//scaninfo") - ] - - nmap_data.update(cls._parse_runstats(root)) - - nmap_data.update(cls._parse_xml_host(root.find(".//host"))) - - return NmapResult.model_validate(nmap_data) - - @classmethod - def _validate_nmap_root(cls, root: ET.Element) -> None: - allowed = { - "scaninfo", - "host", - "runstats", - "verbose", - "debugging", - "taskprogress", - } - - found = {child.tag for child in root} - unexpected = found - allowed - if unexpected: - raise ValueError( - f"Unexpected top-level tags in nmap XML: {sorted(unexpected)}" - ) - - @classmethod - def _parse_scaninfo(cls, scaninfo_el: ET.Element) -> NmapScanInfo: - data = {} - data["type"] = NmapScanType(scaninfo_el.attrib["type"]) - data["protocol"] = IPProtocol(scaninfo_el.attrib["protocol"]) - data["num_services"] = scaninfo_el.attrib["numservices"] - data["services"] = scaninfo_el.attrib["services"] - return NmapScanInfo.model_validate(data) - - @classmethod - def _parse_runstats(cls, root: ET.Element) -> dict: - runstats = root.find("runstats") - if runstats is None: - return {} - - finished = runstats.find("finished") - if finished is None: - return {} - - finished_at = None - ts = finished.attrib.get("time") - if ts: - finished_at = datetime.fromtimestamp(int(ts), tz=UTC) - - return { - "finished_at": finished_at, - "exit_status": finished.attrib.get("exit"), - } - - @classmethod - def _parse_nmaprun(cls, nmaprun_el: ET.Element) -> dict: - nmap_data = {} - nmaprun = dict(nmaprun_el.attrib) - nmap_data["command_line"] = nmaprun["args"] - nmap_data["started_at"] = datetime.fromtimestamp( - float(nmaprun["start"]), tz=UTC - ) - nmap_data["version"] = nmaprun["version"] - nmap_data["xmloutputversion"] = nmaprun["xmloutputversion"] - return nmap_data - - @classmethod - def _parse_xml_host(cls, host_el: ET.Element) -> dict: - """ - Receives a XML tag representing a scanned host with - its services. - """ - data = {} - - # - status_el = host_el.find("status") - data["host_state"] = NmapHostState(status_el.attrib["state"]) - data["host_state_reason"] = NmapHostStatusReason(status_el.attrib["reason"]) - host_state_reason_ttl = status_el.attrib.get("reason_ttl") - if host_state_reason_ttl: - data["host_state_reason_ttl"] = int(host_state_reason_ttl) - - #
- address_el = host_el.find("address") - data["target_ip"] = address_el.attrib["addr"] - - data["hostnames"] = cls._parse_hostnames(host_el.find("hostnames")) - - data["ports"], data["port_stats"] = cls._parse_xml_ports(host_el.find("ports")) - - uptime = host_el.find("uptime") - if uptime is not None: - data["uptime_seconds"] = int(uptime.attrib["seconds"]) - - distance = host_el.find("distance") - if distance is not None: - data["distance"] = int(distance.attrib["value"]) - - tcpsequence = host_el.find("tcpsequence") - if tcpsequence is not None: - data["tcp_sequence_index"] = int(tcpsequence.attrib["index"]) - data["tcp_sequence_difficulty"] = tcpsequence.attrib["difficulty"] - ipidsequence = host_el.find("ipidsequence") - if ipidsequence is not None: - data["ipid_sequence_class"] = ipidsequence.attrib["class"] - tcptssequence = host_el.find("tcptssequence") - if tcptssequence is not None: - data["tcp_timestamp_class"] = tcptssequence.attrib["class"] - - times_elem = host_el.find("times") - if times_elem is not None: - data.update( - { - "srtt_us": int(times_elem.attrib.get("srtt", 0)) or None, - "rttvar_us": int(times_elem.attrib.get("rttvar", 0)) or None, - "timeout_us": int(times_elem.attrib.get("to", 0)) or None, - } - ) - - hostscripts_el = host_el.find("hostscript") - if hostscripts_el is not None: - data["host_scripts"] = [ - NmapHostScript(id=el.attrib["id"], output=el.attrib.get("output")) - for el in hostscripts_el.findall("script") - ] - - data["os_matches"] = cls._parse_os_matches(host_el) - - data["trace"] = cls._parse_trace(host_el) - - return data - - @classmethod - def _parse_os_matches(cls, host_el: ET.Element) -> list[NmapOSMatch] | None: - os_elem = host_el.find("os") - if os_elem is None: - return None - - matches: list[NmapOSMatch] = [] - - for m in os_elem.findall("osmatch"): - classes: list[NmapOSClass] = [] - - for c in m.findall("osclass"): - cpes = [e.text.strip() for e in c.findall("cpe") if e.text] - - classes.append( - NmapOSClass( - vendor=c.attrib.get("vendor"), - osfamily=c.attrib.get("osfamily"), - osgen=c.attrib.get("osgen"), - accuracy=( - int(c.attrib["accuracy"]) - if "accuracy" in c.attrib - else None - ), - cpe=cpes or None, - ) - ) - - matches.append( - NmapOSMatch( - name=m.attrib["name"], - accuracy=int(m.attrib["accuracy"]), - classes=classes, - ) - ) - - return matches or None - - @classmethod - def _parse_hostnames(cls, hostnames_el: ET.Element) -> list[NmapHostname]: - """ - Parses the hostnames element. - e.g. - - - """ - return [ - cls._parse_hostname(hname) for hname in hostnames_el.findall("hostname") - ] - - @classmethod - def _parse_hostname(cls, hostname_el: ET.Element) -> NmapHostname: - """ - Parses the hostname element. - e.g. - - :param hostname_el: XML tag from a nmap scan - """ - return NmapHostname.model_validate(dict(hostname_el.attrib)) - - @classmethod - def _parse_xml_ports( - cls, ports_elem: ET.Element - ) -> tuple[list[NmapPort], NmapPortStats]: - """ - Parses the list of scanned services from a targeted host. - """ - ports: list[NmapPort] = [] - stats = NmapPortStats() - - # handle extraports first - for e in ports_elem.findall("extraports"): - state = PortState(e.attrib["state"]) - count = int(e.attrib["count"]) - - key = state.value.replace("|", "_") - setattr(stats, key, getattr(stats, key) + count) - - for port_elem in ports_elem.findall("port"): - port = cls._parse_xml_port(port_elem) - ports.append(port) - key = port.state.value.replace("|", "_") - setattr(stats, key, getattr(stats, key) + 1) - return ports, stats - - @classmethod - def _parse_xml_service(cls, service_elem: ET.Element) -> NmapService: - svc = { - "name": service_elem.attrib.get("name"), - "product": service_elem.attrib.get("product"), - "version": service_elem.attrib.get("version"), - "extrainfo": service_elem.attrib.get("extrainfo"), - "method": service_elem.attrib.get("method"), - "conf": ( - int(service_elem.attrib["conf"]) - if "conf" in service_elem.attrib - else None - ), - "cpe": [e.text.strip() for e in service_elem.findall("cpe")], - } - - return NmapService.model_validate(svc) - - @classmethod - def _parse_xml_script(cls, script_elem: ET.Element) -> NmapScript: - output = script_elem.attrib.get("output") - if output: - output = output.strip() - script = { - "id": script_elem.attrib["id"], - "output": output, - } - - elements: dict[str, Any] = {} - - # handle value - for elem in script_elem.findall(".//elem"): - key = elem.attrib.get("key") - if key: - elements[key.strip()] = elem.text.strip() - - script["elements"] = elements - return NmapScript.model_validate(script) - - @classmethod - def _parse_xml_port(cls, port_elem: ET.Element) -> NmapPort: - """ - - - - - Username and password - 2 -
- - - """ - state_elem = port_elem.find("state") - - port = { - "port": int(port_elem.attrib["portid"]), - "protocol": port_elem.attrib["protocol"], - "state": PortState(state_elem.attrib["state"]), - "reason": ( - PortStateReason(state_elem.attrib["reason"]) - if "reason" in state_elem.attrib - else None - ), - "reason_ttl": ( - int(state_elem.attrib["reason_ttl"]) - if "reason_ttl" in state_elem.attrib - else None - ), - } - - service_elem = port_elem.find("service") - if service_elem is not None: - port["service"] = cls._parse_xml_service(service_elem) - - port["scripts"] = [] - for script_elem in port_elem.findall("script"): - port["scripts"].append(cls._parse_xml_script(script_elem)) - - return NmapPort.model_validate(port) - - @classmethod - def _parse_trace(cls, host_elem: ET.Element) -> NmapTrace | None: - trace_elem = host_elem.find("trace") - if trace_elem is None: - return None - - port_attr = trace_elem.attrib.get("port") - proto_attr = trace_elem.attrib.get("proto") - - hops: list[NmapTraceHop] = [] - - for hop_elem in trace_elem.findall("hop"): - ttl = hop_elem.attrib.get("ttl") - if ttl is None: - continue # ttl is required by the DTD but guard anyway - - rtt = hop_elem.attrib.get("rtt") - ipaddr = hop_elem.attrib.get("ipaddr") - host = hop_elem.attrib.get("host") - - hops.append( - NmapTraceHop( - ttl=int(ttl), - ipaddr=ipaddr, - rtt_ms=float(rtt) if rtt is not None else None, - host=host, - ) - ) - - return NmapTrace( - port=int(port_attr) if port_attr is not None else None, - protocol=IPProtocol(proto_attr) if proto_attr is not None else None, - hops=hops, - ) - - -def parse_nmap_xml(raw) -> NmapResult: - return NmapXmlParser.parse_xml(raw) diff --git a/generalresearch/models/network/nmap/result.py b/generalresearch/models/network/nmap/result.py deleted file mode 100644 index 57c2e8b..0000000 --- a/generalresearch/models/network/nmap/result.py +++ /dev/null @@ -1,436 +0,0 @@ -from __future__ import annotations - -import json -from datetime import timedelta -from enum import StrEnum -from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal - -from pydantic import BaseModel, Field, computed_field - -from generalresearch.models.network.definitions import IPProtocol - -if TYPE_CHECKING: - from generalresearch.models.custom_types import AwareDatetimeISO, IPvAnyAddressStr - - -class PortState(StrEnum): - OPEN = "open" - CLOSED = "closed" - FILTERED = "filtered" - UNFILTERED = "unfiltered" - OPEN_FILTERED = "open|filtered" - CLOSED_FILTERED = "closed|filtered" - # Added by me, does not get returned. Used for book-keeping - NOT_SCANNED = "not_scanned" - - -class PortStateReason(StrEnum): - SYN_ACK = "syn-ack" - RESET = "reset" - CONN_REFUSED = "conn-refused" - NO_RESPONSE = "no-response" - SYN = "syn" - FIN = "fin" - - ICMP_NET_UNREACH = "net-unreach" - ICMP_HOST_UNREACH = "host-unreach" - ICMP_PROTO_UNREACH = "proto-unreach" - ICMP_PORT_UNREACH = "port-unreach" - - ADMIN_PROHIBITED = "admin-prohibited" - HOST_PROHIBITED = "host-prohibited" - NET_PROHIBITED = "net-prohibited" - - ECHO_REPLY = "echo-reply" - TIME_EXCEEDED = "time-exceeded" - - -class NmapScanType(StrEnum): - SYN = "syn" - CONNECT = "connect" - ACK = "ack" - WINDOW = "window" - MAIMON = "maimon" - FIN = "fin" - NULL = "null" - XMAS = "xmas" - UDP = "udp" - SCTP_INIT = "sctpinit" - SCTP_COOKIE_ECHO = "sctpcookieecho" - - -class NmapHostState(StrEnum): - UP = "up" - DOWN = "down" - UNKNOWN = "unknown" - - -class NmapHostStatusReason(StrEnum): - USER_SET = "user-set" - SYN_ACK = "syn-ack" - RESET = "reset" - ECHO_REPLY = "echo-reply" - ARP_RESPONSE = "arp-response" - NO_RESPONSE = "no-response" - NET_UNREACH = "net-unreach" - HOST_UNREACH = "host-unreach" - PROTO_UNREACH = "proto-unreach" - PORT_UNREACH = "port-unreach" - ADMIN_PROHIBITED = "admin-prohibited" - LOCALHOST_RESPONSE = "localhost-response" - - -class NmapOSClass(BaseModel): - vendor: str = None - osfamily: str = None - osgen: str | None = None - accuracy: int = None - cpe: list[str] | None = None - - -class NmapOSMatch(BaseModel): - name: str - accuracy: int - classes: list[NmapOSClass] = Field(default_factory=list) - - @property - def best_class(self) -> NmapOSClass | None: - if not self.classes: - return None - return max(self.classes, key=lambda m: m.accuracy) - - -class NmapScript(BaseModel): - """ - - """ - - id: str - output: str | None = None - elements: dict[str, Any] = Field(default_factory=dict) - - -class NmapService(BaseModel): - # - name: str | None = None - product: str | None = None - version: str | None = None - extrainfo: str | None = None - method: str | None = None - conf: int | None = None - cpe: list[str] = Field(default_factory=list) - - def model_dump_postgres(self): - d = self.model_dump(mode="json") - d["service_name"] = self.name - return d - - -class NmapPort(BaseModel): - port: int = Field() - protocol: IPProtocol = Field() - # Closed ports will not have a NmapPort record - state: PortState = Field() - reason: PortStateReason | None = Field(default=None) - reason_ttl: int | None = Field(default=None) - - service: NmapService | None = None - scripts: list[NmapScript] = Field(default_factory=list) - - def model_dump_postgres(self, run_id: int): - # Writes for the network_portscanport table - d = {"port_scan_id": run_id} - data = self.model_dump( - mode="json", - include={ - "port", - "state", - "reason", - "reason_ttl", - }, - ) - d.update(data) - d["protocol"] = self.protocol.to_number() - if self.service: - d.update(self.service.model_dump_postgres()) - return d - - -class NmapHostScript(BaseModel): - id: str = Field() - output: str | None = Field(default=None) - - -class NmapTraceHop(BaseModel): - """ - One hop observed during Nmap's traceroute. - - Example XML: - - """ - - ttl: int = Field() - - ipaddr: str | None = Field( - default=None, - description="IP address of the responding router or host", - ) - - rtt_ms: float | None = Field( - default=None, - description="Round-trip time in milliseconds for the probe reaching this hop.", - ) - - host: str | None = Field( - default=None, - description="Reverse DNS hostname for the hop if Nmap resolved one.", - ) - - -class NmapTrace(BaseModel): - """ - Traceroute information collected by Nmap. - - Nmap performs a single traceroute per host using probes matching the scan - type (typically TCP) directed at a chosen destination port. - - Example XML: - - - ... - - """ - - port: int | None = Field( - default=None, - description="Destination port used for traceroute probes (may be absent depending on scan type).", - ) - protocol: IPProtocol | None = Field( - default=None, - description="Transport protocol used for the traceroute probes (tcp, udp, etc.).", - ) - - hops: list[NmapTraceHop] = Field( - default_factory=list, - description="Ordered list of hops observed during the traceroute.", - ) - - @property - def destination(self) -> NmapTraceHop | None: - return self.hops[-1] if self.hops else None - - -class NmapHostname(BaseModel): - # - name: str - type: Literal["PTR", "user"] | None = None - - -class NmapPortStats(BaseModel): - """ - This is counts across all protocols scanned (tcp/udp) - """ - - open: int = 0 - closed: int = 0 - filtered: int = 0 - unfiltered: int = 0 - open_filtered: int = 0 - closed_filtered: int = 0 - - -class NmapScanInfo(BaseModel): - """ - We could have multiple protocols in one run. - - - """ - - type: NmapScanType = Field() - protocol: IPProtocol = Field() - num_services: int = Field() - services: str = Field() - - @cached_property - def port_set(self) -> set[int]: - """ - Expand the Nmap services string into a set of port numbers. - Example: - "22-25,80,443" -> {22,23,24,25,80,443} - """ - ports: set[int] = set() - for part in self.services.split(","): - if "-" in part: - start, end = part.split("-", 1) - ports.update(range(int(start), int(end) + 1)) - else: - ports.add(int(part)) - return ports - - -class NmapResult(BaseModel): - """ - A Nmap Run. Expects that we've only scanned ONE host. - """ - - command_line: str = Field() - started_at: AwareDatetimeISO = Field() - version: str = Field() - xmloutputversion: str = Field() - - scan_infos: list[NmapScanInfo] = Field(min_length=1) - - # comes from - finished_at: AwareDatetimeISO | None = Field(default=None) - exit_status: Literal["success", "error"] | None = Field(default=None) - - ##### - # Everything below here is from within the *single* host we've scanned - ##### - - # - host_state: NmapHostState = Field() - host_state_reason: NmapHostStatusReason = Field() - host_state_reason_ttl: int | None = None - - #
- target_ip: IPvAnyAddressStr = Field() - - hostnames: list[NmapHostname] = Field() - - ports: list[NmapPort] = [] - port_stats: NmapPortStats = Field() - - # - uptime_seconds: int | None = Field(default=None) - # - distance: int | None = Field(description="approx number of hops", default=None) - - # - tcp_sequence_index: int | None = None - tcp_sequence_difficulty: str | None = None - - # - ipid_sequence_class: str | None = None - - # - tcp_timestamp_class: str | None = None - - # - srtt_us: int | None = Field( - default=None, description="smoothed RTT estimate (microseconds µs)" - ) - rttvar_us: int | None = Field( - default=None, description="RTT variance (microseconds µs)" - ) - timeout_us: int | None = Field( - default=None, description="probe timeout (microseconds µs)" - ) - - os_matches: list[NmapOSMatch] | None = Field(default=None) - - host_scripts: list[NmapHostScript] = Field(default_factory=list) - - trace: NmapTrace | None = Field(default=None) - - raw_xml: str | None = None - - @computed_field - @property - def last_boot(self) -> AwareDatetimeISO | None: - if self.uptime_seconds: - return self.started_at - timedelta(seconds=self.uptime_seconds) - - @property - def scan_info_tcp(self): - return next( - filter(lambda x: x.protocol == IPProtocol.TCP, self.scan_infos), None - ) - - @property - def scan_info_udp(self): - return next( - filter(lambda x: x.protocol == IPProtocol.UDP, self.scan_infos), None - ) - - @property - def latency_ms(self) -> float | None: - return self.srtt_us / 1000 if self.srtt_us is not None else None - - @property - def best_os_match(self) -> NmapOSMatch | None: - if not self.os_matches: - return None - return max(self.os_matches, key=lambda m: m.accuracy) - - def filter_ports(self, protocol: IPProtocol, state: PortState) -> list[NmapPort]: - return [p for p in self.ports if p.protocol == protocol and p.state == state] - - @property - def tcp_open_ports(self) -> list[int]: - """ - Returns a list of open TCP port numbers. - """ - return [ - p.port - for p in self.filter_ports(protocol=IPProtocol.TCP, state=PortState.OPEN) - ] - - @property - def udp_open_ports(self) -> list[int]: - """ - Returns a list of open UDP port numbers. - """ - return [ - p.port - for p in self.filter_ports(protocol=IPProtocol.UDP, state=PortState.OPEN) - ] - - @cached_property - def _port_index(self) -> dict[tuple[IPProtocol, int], NmapPort]: - return {(p.protocol, p.port): p for p in self.ports} - - def get_port_state( - self, port: int, protocol: IPProtocol = IPProtocol.TCP - ) -> PortState: - # Explicit (only if scanned and not closed) - if (protocol, port) in self._port_index: - return self._port_index[(protocol, port)].state - - # Check if we even scanned it - scaninfo = next((s for s in self.scan_infos if s.protocol == protocol), None) - if scaninfo and port in scaninfo.port_set: - return PortState.CLOSED - - # We didn't scan it - return PortState.NOT_SCANNED - - def model_dump_postgres(self): - # Writes for the network_portscan table - d = {} - data = self.model_dump( - mode="json", - include={ - "started_at", - "host_state", - "host_state_reason", - "distance", - "uptime_seconds", - "raw_xml", - }, - ) - d.update(data) - d["ip"] = self.target_ip - d["xml_version"] = self.xmloutputversion - d["latency_ms"] = self.latency_ms - d["last_boot"] = self.last_boot - d["parsed"] = self.model_dump_json(indent=0) - d["open_tcp_ports"] = json.dumps(self.tcp_open_ports) - d["open_udp_ports"] = json.dumps(self.udp_open_ports) - return d diff --git a/generalresearch/models/network/rdns/__init__.py b/generalresearch/models/network/rdns/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/generalresearch/models/network/rdns/command.py b/generalresearch/models/network/rdns/command.py deleted file mode 100644 index 2250449..0000000 --- a/generalresearch/models/network/rdns/command.py +++ /dev/null @@ -1,38 +0,0 @@ -import subprocess -from typing import TYPE_CHECKING - -from generalresearch.models.network.rdns.parser import parse_rdns_output - -if TYPE_CHECKING: - from generalresearch.models.network.rdns.result import RDNSResult - from generalresearch.models.network.tool_run_command import RDNSRunCommand - - -def run_rdns(config: RDNSRunCommand) -> RDNSResult: - cmd = config.to_command_str() - args = cmd.split(" ") - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - return parse_rdns_output(ip=config.options.ip, raw=raw) - - -def build_rdns_command(ip: str) -> str: - # e.g. dig +noall +answer -x 1.2.3.4 - return f"dig +noall +answer -x {ip}" - - -def get_dig_version() -> str: - proc = subprocess.run( - ["dig", "-v"], - capture_output=True, - text=True, - check=False, - ) - # e.g. DiG 9.18.39-0ubuntu0.22.04.2-Ubuntu - ver_str = proc.stderr.strip() + proc.stdout.strip() - return ver_str.split("-", 1)[0].split(" ", 1)[1] diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py deleted file mode 100644 index d6de84b..0000000 --- a/generalresearch/models/network/rdns/execute.py +++ /dev/null @@ -1,44 +0,0 @@ -from __future__ import annotations - -from datetime import UTC, datetime -from uuid import uuid4 - -from generalresearch.models.custom_types import UUIDStr -from generalresearch.models.network.rdns.command import ( - get_dig_version, - run_rdns, -) -from generalresearch.models.network.tool_run import ( - RDNSRun, - Status, - ToolClass, - ToolName, -) -from generalresearch.models.network.tool_run_command import ( - RDNSRunCommand, - RDNSRunCommandOptions, -) - - -def execute_rdns(ip: str, scan_group_id: UUIDStr | None = None): - started_at = datetime.now(tz=UTC) - tool_version = get_dig_version() - config = RDNSRunCommand(options=RDNSRunCommandOptions(ip=ip)) - result = run_rdns(config) - finished_at = datetime.now(tz=UTC) - - run = RDNSRun( - tool_name=ToolName.DIG, - tool_class=ToolClass.RDNS, - tool_version=tool_version, - status=Status.SUCCESS, - ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=config.to_command_str(), - scan_group_id=scan_group_id or uuid4().hex, - config=config, - parsed=result, - ) - - return run diff --git a/generalresearch/models/network/rdns/parser.py b/generalresearch/models/network/rdns/parser.py deleted file mode 100644 index 31a5ed6..0000000 --- a/generalresearch/models/network/rdns/parser.py +++ /dev/null @@ -1,21 +0,0 @@ -import ipaddress -import re - -from generalresearch.models.custom_types import IPvAnyAddressStr -from generalresearch.models.network.rdns.result import RDNSResult - -PTR_RE = re.compile(r"\sPTR\s+([^\s]+)\.") - - -def parse_rdns_output(ip: IPvAnyAddressStr, raw: str) -> RDNSResult: - hostnames: list[str] = [] - - for line in raw.splitlines(): - m = PTR_RE.search(line) - if m: - hostnames.append(m.group(1)) - - return RDNSResult( - ip=ipaddress.ip_address(ip), - hostnames=hostnames, - ) diff --git a/generalresearch/models/network/rdns/result.py b/generalresearch/models/network/rdns/result.py deleted file mode 100644 index 46af643..0000000 --- a/generalresearch/models/network/rdns/result.py +++ /dev/null @@ -1,52 +0,0 @@ -from __future__ import annotations - -import json -from functools import cached_property - -import tldextract -from pydantic import BaseModel, Field, computed_field, model_validator - -from generalresearch.models.custom_types import IPvAnyAddressStr - - -class RDNSResult(BaseModel): - - ip: IPvAnyAddressStr = Field() - - hostnames: list[str] = Field(default_factory=list) - - @model_validator(mode="after") - def validate_hostname_prop(self): - assert len(self.hostnames) == self.hostname_count - if self.hostnames: - assert self.hostnames[0] == self.primary_hostname - assert self.primary_domain in self.primary_hostname - return self - - @computed_field(examples=["fixed-187-191-8-145.totalplay.net"]) - @cached_property - def primary_hostname(self) -> str | None: - if self.hostnames: - return self.hostnames[0] - - @computed_field(examples=[1]) - @cached_property - def hostname_count(self) -> int: - return len(self.hostnames) - - @computed_field(examples=["totalplay.net"]) - @cached_property - def primary_domain(self) -> str | None: - if self.primary_hostname: - return tldextract.extract( - self.primary_hostname - ).top_domain_under_public_suffix - - def model_dump_postgres(self): - # Writes for the network_rdnsresult table - d = self.model_dump( - mode="json", - include={"primary_hostname", "primary_domain", "hostname_count"}, - ) - d["hostnames"] = json.dumps(self.hostnames) - return d diff --git a/generalresearch/models/network/tool_run.py b/generalresearch/models/network/tool_run.py deleted file mode 100644 index 9088fe3..0000000 --- a/generalresearch/models/network/tool_run.py +++ /dev/null @@ -1,121 +0,0 @@ -from __future__ import annotations - -from enum import StrEnum -from typing import TYPE_CHECKING, Literal -from uuid import uuid4 - -from pydantic import BaseModel, Field, PositiveInt - -from generalresearch.models.custom_types import ( - AwareDatetimeISO, - IPvAnyAddressStr, - UUIDStr, -) - -if TYPE_CHECKING: - - from generalresearch.models.network.mtr.result import MTRResult - from generalresearch.models.network.nmap.result import NmapResult - from generalresearch.models.network.rdns.result import RDNSResult - from generalresearch.models.network.tool_run_command import ( - MTRRunCommand, - NmapRunCommand, - RDNSRunCommand, - ToolRunCommand, - ) - - -class ToolClass(StrEnum): - PORT_SCAN = "port_scan" - RDNS = "rdns" - PING = "ping" - TRACEROUTE = "traceroute" - - -class ToolName(StrEnum): - NMAP = "nmap" - RUSTMAP = "rustmap" - DIG = "dig" - PING = "ping" - TRACEROUTE = "traceroute" - MTR = "mtr" - - -class Status(StrEnum): - SUCCESS = "success" - FAILED = "failed" - TIMEOUT = "timeout" - ERROR = "error" - - -class ToolRun(BaseModel): - """ - A run of a networking tool against one host/ip. - """ - - id: PositiveInt | None = Field(default=None) - - ip: IPvAnyAddressStr = Field() - scan_group_id: UUIDStr = Field(default_factory=lambda: uuid4().hex) - tool_class: ToolClass = Field() - tool_name: ToolName = Field() - tool_version: str = Field() - - started_at: AwareDatetimeISO = Field() - finished_at: AwareDatetimeISO | None = Field(default=None) - status: Status | None = Field(default=None) - - raw_command: str = Field() - - config: ToolRunCommand = Field() - - def model_dump_postgres(self): - d = self.model_dump(mode="json", exclude={"config"}) - d["config"] = self.config.model_dump_json() - return d - - -class NmapRun(ToolRun): - tool_class: Literal[ToolClass.PORT_SCAN] = Field(default=ToolClass.PORT_SCAN) - tool_name: Literal[ToolName.NMAP] = Field(default=ToolName.NMAP) - config: NmapRunCommand = Field() - - parsed: NmapResult = Field() - - def model_dump_postgres(self): - d = super().model_dump_postgres() - d["run_id"] = self.id - d.update(self.parsed.model_dump_postgres()) - return d - - -class RDNSRun(ToolRun): - tool_class: Literal[ToolClass.RDNS] = Field(default=ToolClass.RDNS) - tool_name: Literal[ToolName.DIG] = Field(default=ToolName.DIG) - config: RDNSRunCommand = Field() - - parsed: RDNSResult = Field() - - def model_dump_postgres(self): - d = super().model_dump_postgres() - d["run_id"] = self.id - d.update(self.parsed.model_dump_postgres()) - return d - - -class MTRRun(ToolRun): - tool_class: Literal[ToolClass.TRACEROUTE] = Field(default=ToolClass.TRACEROUTE) - tool_name: Literal[ToolName.MTR] = Field(default=ToolName.MTR) - config: MTRRunCommand = Field() - - facility_id: int = Field(default=1) - source_ip: IPvAnyAddressStr = Field() - parsed: MTRResult = Field() - - def model_dump_postgres(self): - d = super().model_dump_postgres() - d["run_id"] = self.id - d["source_ip"] = self.source_ip - d["facility_id"] = self.facility_id - d.update(self.parsed.model_dump_postgres()) - return d diff --git a/generalresearch/models/network/tool_run_command.py b/generalresearch/models/network/tool_run_command.py deleted file mode 100644 index b07b811..0000000 --- a/generalresearch/models/network/tool_run_command.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Literal - -from pydantic import BaseModel, Field - -from generalresearch.models.network.definitions import IPProtocol - -if TYPE_CHECKING: - from generalresearch.models.custom_types import IPvAnyAddressStr - - -class ToolRunCommand(BaseModel): - command: str = Field() - options: dict[str, str | int | None] = Field(default_factory=dict) - - -class NmapRunCommandOptions(BaseModel): - ip: IPvAnyAddressStr - top_ports: int | None = Field(default=1000) - ports: str | None = Field(default=None) - no_ping: bool = Field(default=True) - enable_advanced: bool = Field(default=True) - timing: int = Field(default=4) - - -class NmapRunCommand(ToolRunCommand): - command: Literal["nmap"] = Field(default="nmap") - options: NmapRunCommandOptions = Field() - - def to_command_str(self): - from generalresearch.models.network.nmap.command import build_nmap_command - - options = self.options - return build_nmap_command(**options.model_dump()) - - -class RDNSRunCommandOptions(BaseModel): - ip: IPvAnyAddressStr - - -class RDNSRunCommand(ToolRunCommand): - command: Literal["dig"] = Field(default="dig") - options: RDNSRunCommandOptions = Field() - - def to_command_str(self): - from generalresearch.models.network.rdns.command import build_rdns_command - - options = self.options - return build_rdns_command(**options.model_dump()) - - -class MTRRunCommandOptions(BaseModel): - ip: IPvAnyAddressStr = Field() - protocol: IPProtocol = Field(default=IPProtocol.ICMP) - port: int | None = Field(default=None) - report_cycles: int = Field(default=10) - - -class MTRRunCommand(ToolRunCommand): - command: Literal["mtr"] = Field(default="mtr") - options: MTRRunCommandOptions = Field() - - def to_command_str(self): - from generalresearch.models.network.mtr.command import build_mtr_command - - options = self.options - return build_mtr_command(**options.model_dump()) diff --git a/generalresearch/models/network/utils.py b/generalresearch/models/network/utils.py deleted file mode 100644 index fee9b80..0000000 --- a/generalresearch/models/network/utils.py +++ /dev/null @@ -1,5 +0,0 @@ -import requests - - -def get_source_ip(): - return requests.get("https://icanhazip.com?").text.strip() diff --git a/generalresearch/models/pollfish/question.py b/generalresearch/models/pollfish/question.py index f0c733c..aadf71e 100644 --- a/generalresearch/models/pollfish/question.py +++ b/generalresearch/models/pollfish/question.py @@ -4,17 +4,15 @@ from __future__ import annotations import json import logging from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from pydantic import BaseModel, Field, model_validator from generalresearch.models.definitions import Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion - -if TYPE_CHECKING: - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/precision/question.py b/generalresearch/models/precision/question.py index 3e39124..97f6ca1 100644 --- a/generalresearch/models/precision/question.py +++ b/generalresearch/models/precision/question.py @@ -4,22 +4,20 @@ from __future__ import annotations import json import logging from enum import StrEnum -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator from generalresearch.models.definitions import Source +from generalresearch.models.precision import PrecisionQuestionID from generalresearch.models.string_utils import remove_nbsp from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, MarketplaceUserQuestionAnswer, ) - -if TYPE_CHECKING: - from generalresearch.models.precision import PrecisionQuestionID - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/precision/survey.py b/generalresearch/models/precision/survey.py index b77a365..cebe155 100644 --- a/generalresearch/models/precision/survey.py +++ b/generalresearch/models/precision/survey.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from datetime import UTC from functools import cached_property -from typing import TYPE_CHECKING, Annotated, Any, Literal, Self +from typing import Annotated, Any, Literal, Self from more_itertools import flatten from pydantic import ( @@ -23,7 +23,7 @@ from generalresearch.models.custom_types import ( UUIDStrCoerce, ) from generalresearch.models.definitions import Source -from generalresearch.models.precision import PrecisionStatus +from generalresearch.models.precision import PrecisionQuestionID, PrecisionStatus from generalresearch.models.thl.demographics import Gender from generalresearch.models.thl.survey import MarketplaceTask from generalresearch.models.thl.survey.condition import ( @@ -31,9 +31,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - from generalresearch.models.precision import PrecisionQuestionID - class PrecisionCondition(MarketplaceCondition): question_id: PrecisionQuestionID | None = Field() diff --git a/generalresearch/models/precision/task_collection.py b/generalresearch/models/precision/task_collection.py index 71241ad..c8db2af 100644 --- a/generalresearch/models/precision/task_collection.py +++ b/generalresearch/models/precision/task_collection.py @@ -1,18 +1,16 @@ -from typing import TYPE_CHECKING, Any +from typing import Any import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.precision import PrecisionStatus +from generalresearch.models.precision.survey import PrecisionSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.precision.survey import PrecisionSurvey - COUNTRY_ISOS = Localelator().get_all_countries() LANGUAGE_ISOS = Localelator().get_all_languages() diff --git a/generalresearch/models/prodege/question.py b/generalresearch/models/prodege/question.py index b963785..c0160bd 100644 --- a/generalresearch/models/prodege/question.py +++ b/generalresearch/models/prodege/question.py @@ -6,7 +6,7 @@ import logging from datetime import UTC, datetime from enum import StrEnum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from pydantic import ( BaseModel, @@ -18,15 +18,13 @@ from pydantic import ( ) from generalresearch.locales import Localelator +from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.definitions import MAX_INT32, Source +from generalresearch.models.prodege import ProdegeQuestionIdType from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion - -if TYPE_CHECKING: - from generalresearch.models.custom_types import AwareDatetimeISO - from generalresearch.models.prodege import ProdegeQuestionIdType - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/prodege/survey.py b/generalresearch/models/prodege/survey.py index 26898d0..27034b0 100644 --- a/generalresearch/models/prodege/survey.py +++ b/generalresearch/models/prodege/survey.py @@ -7,7 +7,7 @@ from collections import defaultdict from datetime import UTC, datetime from decimal import Decimal from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from pydantic import ( BaseModel, @@ -34,7 +34,9 @@ from generalresearch.models.definitions import ( ) from generalresearch.models.prodege import ( ProdegePastParticipationType, + ProdegeQuestionIdType, ProdegeStatus, + ProdgeRedirectStatus, ) from generalresearch.models.prodege.definitions import PG_COUNTRY_TO_ISO from generalresearch.models.thl.demographics import Gender @@ -44,13 +46,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - - from generalresearch.models.prodege import ( - ProdegeQuestionIdType, - ProdgeRedirectStatus, - ) - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) diff --git a/generalresearch/models/prodege/task_collection.py b/generalresearch/models/prodege/task_collection.py index 774fc7b..9f6a81b 100644 --- a/generalresearch/models/prodege/task_collection.py +++ b/generalresearch/models/prodege/task_collection.py @@ -1,18 +1,16 @@ -from typing import TYPE_CHECKING, Any +from typing import Any import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.prodege import ProdegeStatus +from generalresearch.models.prodege.survey import ProdegeSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.prodege.survey import ProdegeSurvey - COUNTRY_ISOS = Localelator().get_all_countries() LANGUAGE_ISOS = Localelator().get_all_languages() diff --git a/generalresearch/models/repdata/question.py b/generalresearch/models/repdata/question.py index 0ec102b..5115426 100644 --- a/generalresearch/models/repdata/question.py +++ b/generalresearch/models/repdata/question.py @@ -4,7 +4,7 @@ import json import logging from enum import StrEnum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from uuid import UUID from pydantic import ( @@ -20,11 +20,9 @@ from pydantic import ( from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion - -if TYPE_CHECKING: - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/repdata/task_collection.py b/generalresearch/models/repdata/task_collection.py index f2cb63b..aa591a9 100644 --- a/generalresearch/models/repdata/task_collection.py +++ b/generalresearch/models/repdata/task_collection.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index @@ -8,14 +8,12 @@ from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.repdata import RepDataStatus +from generalresearch.models.repdata.survey import RepDataSurveyHashed from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.repdata.survey import RepDataSurveyHashed - COUNTRY_ISOS = Localelator().get_all_countries() LANGUAGE_ISOS = Localelator().get_all_languages() diff --git a/generalresearch/models/sago/question.py b/generalresearch/models/sago/question.py index 216b278..148b015 100644 --- a/generalresearch/models/sago/question.py +++ b/generalresearch/models/sago/question.py @@ -6,7 +6,7 @@ import json import logging from enum import StrEnum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from pydantic import ( BaseModel, @@ -18,15 +18,13 @@ from pydantic import ( model_validator, ) +from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.definitions import MAX_INT32, Source from generalresearch.models.string_utils import remove_nbsp from generalresearch.models.thl.profiling.marketplace import MarketplaceQuestion - -if TYPE_CHECKING: - from generalresearch.models.custom_types import AwareDatetimeISO - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/sago/survey.py b/generalresearch/models/sago/survey.py index c9bf431..5330ace 100644 --- a/generalresearch/models/sago/survey.py +++ b/generalresearch/models/sago/survey.py @@ -5,7 +5,7 @@ import logging from datetime import UTC from decimal import Decimal from functools import cached_property -from typing import TYPE_CHECKING, Annotated, Any, Literal, Self +from typing import Annotated, Any, Literal, Self from more_itertools import flatten from pydantic import ( @@ -18,6 +18,14 @@ from pydantic import ( ) from generalresearch.locales import Localelator +from generalresearch.models.custom_types import ( + AlphaNumStr, + AlphaNumStrSet, + AwareDatetimeISO, + CoercedStr, + DeviceTypes, + IPLikeStrSet, +) from generalresearch.models.definitions import LogicalOperator, Source from generalresearch.models.sago import SagoStatus from generalresearch.models.thl.demographics import Gender @@ -27,16 +35,6 @@ from generalresearch.models.thl.survey.condition import ( MarketplaceCondition, ) -if TYPE_CHECKING: - from generalresearch.models.custom_types import ( - AlphaNumStr, - AlphaNumStrSet, - AwareDatetimeISO, - CoercedStr, - DeviceTypes, - IPLikeStrSet, - ) - logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) diff --git a/generalresearch/models/sago/task_collection.py b/generalresearch/models/sago/task_collection.py index 490f7a0..2879d9c 100644 --- a/generalresearch/models/sago/task_collection.py +++ b/generalresearch/models/sago/task_collection.py @@ -1,20 +1,18 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.sago import SagoStatus +from generalresearch.models.sago.survey import SagoSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.sago.survey import SagoSurvey - COUNTRY_ISOS: set[str] = Localelator().get_all_countries() LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() diff --git a/generalresearch/models/spectrum/question.py b/generalresearch/models/spectrum/question.py index 4b854fb..839f7a8 100644 --- a/generalresearch/models/spectrum/question.py +++ b/generalresearch/models/spectrum/question.py @@ -6,7 +6,7 @@ import logging from datetime import UTC, datetime from enum import IntEnum, StrEnum from functools import cached_property -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from uuid import UUID from pydantic import ( @@ -18,18 +18,16 @@ from pydantic import ( model_validator, ) +from generalresearch.models.custom_types import AwareDatetimeISO from generalresearch.models.definitions import MAX_INT32, Source +from generalresearch.models.spectrum import SpectrumQuestionIdType from generalresearch.models.string_utils import remove_nbsp from generalresearch.models.thl.profiling.marketplace import ( MarketplaceQuestion, ) - -if TYPE_CHECKING: - from generalresearch.models.custom_types import AwareDatetimeISO - from generalresearch.models.spectrum import SpectrumQuestionIdType - from generalresearch.models.thl.profiling.upk_question import ( - UpkQuestion, - ) +from generalresearch.models.thl.profiling.upk_question import ( + UpkQuestion, +) logging.basicConfig() logger = logging.getLogger() diff --git a/generalresearch/models/spectrum/survey.py b/generalresearch/models/spectrum/survey.py index 6689d27..5b330a8 100644 --- a/generalresearch/models/spectrum/survey.py +++ b/generalresearch/models/spectrum/survey.py @@ -55,7 +55,7 @@ class SpectrumCondition(MarketplaceCondition): try: values = [tuple(map(int, v.split("-"))) for v in self.values] assert all(len(x) == 2 for x in values) - except (ValueError, AssertionError): + except ValueError, AssertionError: return self self.values = sorted( {str(val) for tupl in values for val in range(tupl[0], tupl[1] + 1)} @@ -75,7 +75,7 @@ class SpectrumCondition(MarketplaceCondition): rs["from"] = round(rs["from"] / 12) rs["to"] = round(rs["to"] / 12) d["values"] = [ - f"{rs["from"] or "inf"}-{rs["to"] or "inf"}" for rs in d["range_sets"] + f"{rs['from'] or 'inf'}-{rs['to'] or 'inf'}" for rs in d["range_sets"] ] d["value_type"] = ConditionValueType.RANGE return cls.model_validate(d) diff --git a/generalresearch/models/spectrum/task_collection.py b/generalresearch/models/spectrum/task_collection.py index 8e49434..609114e 100644 --- a/generalresearch/models/spectrum/task_collection.py +++ b/generalresearch/models/spectrum/task_collection.py @@ -1,21 +1,17 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import pandas as pd from pandera.pandas import Check, Column, DataFrameSchema, Index from generalresearch.locales import Localelator from generalresearch.models.definitions import TaskCalculationType from generalresearch.models.spectrum import SpectrumStatus +from generalresearch.models.spectrum.survey import SpectrumSurvey from generalresearch.models.thl.survey.task_collection import ( TaskCollection, create_empty_df_from_schema, ) -if TYPE_CHECKING: - from generalresearch.models.spectrum.survey import SpectrumSurvey - COUNTRY_ISOS: set[str] = Localelator().get_all_countries() LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() diff --git a/generalresearch/models/thl/contest/__init__.py b/generalresearch/models/thl/contest/__init__.py index c8342b3..0d7ace5 100644 --- a/generalresearch/models/thl/contest/__init__.py +++ b/generalresearch/models/thl/contest/__init__.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any, Self +from typing import Any, Self from uuid import uuid4 from pydantic import ( @@ -12,12 +12,10 @@ from pydantic import ( model_validator, ) +from generalresearch.currency import USDCent from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.thl.contest.definitions import ContestPrizeKind - -if TYPE_CHECKING: - from generalresearch.currency import USDCent - from generalresearch.models.thl.user import User +from generalresearch.models.thl.user import User class ContestEntryRule(BaseModel): @@ -88,9 +86,9 @@ class ContestPrize(BaseModel): @model_validator(mode="after") def validate_cash_value(self) -> Self: if self.kind == ContestPrizeKind.CASH: - assert ( - self.estimated_cash_value == self.cash_amount - ), "if kind is CASH, cash_amount must equal estimated_cash_value" + assert self.estimated_cash_value == self.cash_amount, ( + "if kind is CASH, cash_amount must equal estimated_cash_value" + ) return self diff --git a/generalresearch/models/thl/contest/contest.py b/generalresearch/models/thl/contest/contest.py index 5e30778..6fc60f6 100644 --- a/generalresearch/models/thl/contest/contest.py +++ b/generalresearch/models/thl/contest/contest.py @@ -3,7 +3,7 @@ from __future__ import annotations import json from abc import ABC, abstractmethod from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any, Self +from typing import Any, Self from uuid import uuid4 from pydantic import ( @@ -19,18 +19,14 @@ from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.thl.contest import ( ContestEndCondition, ContestPrize, + ContestWinner, ) from generalresearch.models.thl.contest.definitions import ( ContestEndReason, ContestStatus, ContestType, ) - -if TYPE_CHECKING: - from generalresearch.models.thl.contest import ( - ContestWinner, - ) - from generalresearch.models.thl.locales import CountryISOs +from generalresearch.models.thl.locales import CountryISOs class ContestBase(BaseModel, ABC): diff --git a/generalresearch/models/thl/contest/contest_entry.py b/generalresearch/models/thl/contest/contest_entry.py index 261b3fc..4e90eb5 100644 --- a/generalresearch/models/thl/contest/contest_entry.py +++ b/generalresearch/models/thl/contest/contest_entry.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any +from typing import Any from uuid import uuid4 from pydantic import ( @@ -16,9 +16,7 @@ from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr from generalresearch.models.thl.contest.definitions import ( ContestEntryType, ) - -if TYPE_CHECKING: - from generalresearch.models.thl.user import User +from generalresearch.models.thl.user import User class ContestEntryCreate(BaseModel): diff --git a/generalresearch/models/thl/contest/leaderboard.py b/generalresearch/models/thl/contest/leaderboard.py index efbbd0a..064c0d1 100644 --- a/generalresearch/models/thl/contest/leaderboard.py +++ b/generalresearch/models/thl/contest/leaderboard.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from pydantic import ( ConfigDict, @@ -16,6 +16,9 @@ from generalresearch.currency import USDCent from generalresearch.decorators import LOG from generalresearch.managers.leaderboard import country_timezone from generalresearch.managers.leaderboard.manager import LeaderboardManager +from generalresearch.managers.thl.user_manager.user_manager import ( + UserManager, +) from generalresearch.models.thl.contest import ( ContestEndCondition, ContestPrize, @@ -39,11 +42,6 @@ from generalresearch.models.thl.leaderboard import ( LeaderboardFrequency, ) -if TYPE_CHECKING: - from generalresearch.managers.thl.user_manager.user_manager import ( - UserManager, - ) - class LeaderboardContestCreate(ContestBase): model_config = ConfigDict( @@ -75,9 +73,9 @@ class LeaderboardContestCreate(ContestBase): ranks = {x.leaderboard_rank for x in self.prizes} assert None not in ranks, "Must have leaderboard_rank defined" assert min(ranks) == 1, "Must start with rank 1" - assert ranks == set( - range(min(ranks), max(ranks) + 1) - ), "cannot skip prize leaderboard_ranks" + assert ranks == set(range(min(ranks), max(ranks) + 1)), ( + "cannot skip prize leaderboard_ranks" + ) return self @model_validator(mode="after") @@ -88,9 +86,9 @@ class LeaderboardContestCreate(ContestBase): @model_validator(mode="after") def check_end_condition(self) -> Self: - assert ( - not self.end_condition.target_entry_amount - ), "target_entry_amount not valid in leaderboard contest" + assert not self.end_condition.target_entry_amount, ( + "target_entry_amount not valid in leaderboard contest" + ) # the ends_at will get set automatically from the leaderboard_key return self @@ -174,13 +172,13 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): @model_validator(mode="after") def validate_product_lb_key(self) -> Self: - assert ( - self.product_id == self.leaderboard_key_parts["product_id"] - ), "leaderboard_key product_id is invalid" + assert self.product_id == self.leaderboard_key_parts["product_id"], ( + "leaderboard_key product_id is invalid" + ) if self.country_isos: - assert ( - len(self.country_isos) == 1 - ), "Can only set 1 country_iso in a leaderboard contest" + assert len(self.country_isos) == 1, ( + "Can only set 1 country_iso in a leaderboard contest" + ) assert ( next(iter(self.country_isos)) == self.leaderboard_key_parts["country_iso"] @@ -192,9 +190,9 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): @model_validator(mode="after") def validate_tie_break(self) -> Self: if self.tie_break_strategy == LeaderboardTieBreakStrategy.SPLIT_PRIZE_POOL: - assert all( - p.kind == ContestPrizeKind.CASH for p in self.prizes - ), "All prizes must be cash due to the tie-break strategy" + assert all(p.kind == ContestPrizeKind.CASH for p in self.prizes), ( + "All prizes must be cash due to the tie-break strategy" + ) return self @model_validator(mode="after") diff --git a/generalresearch/models/thl/contest/raffle.py b/generalresearch/models/thl/contest/raffle.py index 21bc481..7e84e89 100644 --- a/generalresearch/models/thl/contest/raffle.py +++ b/generalresearch/models/thl/contest/raffle.py @@ -4,7 +4,7 @@ import logging import random from collections import defaultdict from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any, Literal, Self +from typing import Any, Literal, Self from pydantic import ( ConfigDict, @@ -20,12 +20,16 @@ from generalresearch.models.thl.contest import ( ContestEndCondition, ContestEntryRule, ContestPrize, + ContestWinner, ) from generalresearch.models.thl.contest.contest import ( Contest, ContestBase, ContestUserView, ) +from generalresearch.models.thl.contest.contest_entry import ( + ContestEntry, +) from generalresearch.models.thl.contest.definitions import ( ContestEndReason, ContestEntryType, @@ -34,14 +38,6 @@ from generalresearch.models.thl.contest.definitions import ( ContestType, ) -if TYPE_CHECKING: - from generalresearch.models.thl.contest import ( - ContestWinner, - ) - from generalresearch.models.thl.contest.contest_entry import ( - ContestEntry, - ) - logging.basicConfig() LOG = logging.getLogger() LOG.setLevel(logging.INFO) @@ -114,9 +110,9 @@ class RaffleContest(RaffleContestCreate, Contest): @model_validator(mode="after") def validate_entry_type(self): - assert all( - entry.entry_type == self.entry_type for entry in self.entries - ), f"all entries must be of type {self.entry_type}" + assert all(entry.entry_type == self.entry_type for entry in self.entries), ( + f"all entries must be of type {self.entry_type}" + ) return self @field_validator("current_amount", mode="before") diff --git a/generalresearch/models/thl/profiling/upk_property.py b/generalresearch/models/thl/profiling/upk_property.py index 96f1b4c..e46e00a 100644 --- a/generalresearch/models/thl/profiling/upk_property.py +++ b/generalresearch/models/thl/profiling/upk_property.py @@ -2,17 +2,14 @@ from __future__ import annotations from enum import StrEnum from functools import cached_property -from typing import TYPE_CHECKING from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from generalresearch.models.custom_types import CountryISOLike, UUIDStr +from generalresearch.models.thl.category import Category from generalresearch.utils.enum import ReprEnumMeta -if TYPE_CHECKING: - from generalresearch.models.thl.category import Category - class PropertyType(StrEnum, metaclass=ReprEnumMeta): # UserProfileKnowledge Item diff --git a/generalresearch/models/thl/profiling/user_info.py b/generalresearch/models/thl/profiling/user_info.py index 5124d17..46733d4 100644 --- a/generalresearch/models/thl/profiling/user_info.py +++ b/generalresearch/models/thl/profiling/user_info.py @@ -1,18 +1,14 @@ from __future__ import annotations -from typing import TYPE_CHECKING - from pydantic import BaseModel, ConfigDict, Field from pydantic.json_schema import SkipJsonSchema from generalresearch.models.custom_types import AwareDatetimeISO - -if TYPE_CHECKING: - from generalresearch.models.definitions import Source - from generalresearch.models.thl.profiling.user_question_answer import ( - MarketplaceResearchProfileQuestion, - ) - from generalresearch.models.thl.user import User +from generalresearch.models.definitions import Source +from generalresearch.models.thl.profiling.user_question_answer import ( + MarketplaceResearchProfileQuestion, +) +from generalresearch.models.thl.user import User class UserProfileKnowledgeAnswer(BaseModel): diff --git a/generalresearch/models/thl/survey/__init__.py b/generalresearch/models/thl/survey/__init__.py index 76f819e..7749bdb 100644 --- a/generalresearch/models/thl/survey/__init__.py +++ b/generalresearch/models/thl/survey/__init__.py @@ -3,32 +3,27 @@ from __future__ import annotations from abc import ABC, abstractmethod from decimal import Decimal from itertools import product -from typing import TYPE_CHECKING from more_itertools import flatten from pydantic import BaseModel, Field +from generalresearch.models.definitions import Source from generalresearch.models.thl.demographics import ( AgeGroup, DemographicTarget, Gender, ) +from generalresearch.models.thl.locales import ( + CountryISO, + CountryISOs, + LanguageISO, + LanguageISOs, +) from generalresearch.models.thl.survey.condition import ( ConditionValueType, + MarketplaceCondition, ) -if TYPE_CHECKING: - from generalresearch.models.definitions import Source - from generalresearch.models.thl.locales import ( - CountryISO, - CountryISOs, - LanguageISO, - LanguageISOs, - ) - from generalresearch.models.thl.survey.condition import ( - MarketplaceCondition, - ) - class MarketplaceTask(BaseModel, ABC): """This is called a "Task" even though generally it represents a survey diff --git a/test_utils/managers/network/__init__.py b/test_utils/managers/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/test_utils/managers/network/conftest.py b/test_utils/managers/network/conftest.py deleted file mode 100644 index e69de29..0000000 diff --git a/test_utils/models/network/__init__.py b/test_utils/models/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/test_utils/models/network/conftest.py b/test_utils/models/network/conftest.py deleted file mode 100644 index c7bcc7e..0000000 --- a/test_utils/models/network/conftest.py +++ /dev/null @@ -1,145 +0,0 @@ -import os -from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING -from uuid import uuid4 - -import pytest -from pytest import FixtureRequest as Request - -from generalresearch.managers.network.label import IPLabelManager -from generalresearch.managers.network.tool_run import ToolRunManager -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.mtr.parser import parse_mtr_output -from generalresearch.models.network.mtr.result import MTRResult -from generalresearch.models.network.nmap.parser import parse_nmap_xml -from generalresearch.models.network.nmap.result import NmapResult -from generalresearch.models.network.rdns.parser import parse_rdns_output -from generalresearch.models.network.rdns.result import RDNSResult -from generalresearch.models.network.tool_run import MTRRun, NmapRun, RDNSRun, Status -from generalresearch.models.network.tool_run_command import ( - MTRRunCommand, - MTRRunCommandOptions, - NmapRunCommand, - NmapRunCommandOptions, - RDNSRunCommand, - RDNSRunCommandOptions, -) -from generalresearch.pg_helper import PostgresConfig - - -@pytest.fixture(scope="session") -def scan_group_id() -> str: - return uuid4().hex - - -@pytest.fixture(scope="session") -def iplabel_manager(thl_web_rw: PostgresConfig) -> IPLabelManager: - return IPLabelManager(pg_config=thl_web_rw) - - -@pytest.fixture(scope="session") -def toolrun_manager(thl_web_rw: PostgresConfig) -> ToolRunManager: - return ToolRunManager(pg_config=thl_web_rw) - - -@pytest.fixture(scope="session") -def nmap_raw_output(request: Request) -> str: - fp = os.path.join(request.config.rootpath, "data/nmaprun1.xml") - with open(fp) as f: - data = f.read() - return data - - -@pytest.fixture(scope="session") -def nmap_result(nmap_raw_output: str) -> NmapResult: - return parse_nmap_xml(nmap_raw_output) - - -@pytest.fixture(scope="session") -def nmap_run(nmap_result: NmapResult, scan_group_id: str): - r = nmap_result - config = NmapRunCommand( - command="nmap", - options=NmapRunCommandOptions( - ip=r.target_ip, ports="22-1000,11000,1100,3389,61232", top_ports=None - ), - ) - return NmapRun( - tool_version=r.version, - status=Status.SUCCESS, - ip=r.target_ip, - started_at=r.started_at, - finished_at=r.finished_at, - raw_command=config.to_command_str(), - scan_group_id=scan_group_id, - config=config, - parsed=r, - ) - - -@pytest.fixture(scope="session") -def dig_raw_output() -> str: - return "156.32.33.45.in-addr.arpa. 300 IN PTR scanme.nmap.org." - - -@pytest.fixture(scope="session") -def rdns_result(dig_raw_output: str) -> RDNSResult: - return parse_rdns_output(ip="45.33.32.156", raw=dig_raw_output) - - -@pytest.fixture(scope="session") -def rdns_run(rdns_result: RDNSResult, scan_group_id: str): - r = rdns_result - ip = "45.33.32.156" - utc_now = datetime.now(tz=UTC) - config = RDNSRunCommand(command="dig", options=RDNSRunCommandOptions(ip=ip)) - return RDNSRun( - tool_version="1.2.3", - status=Status.SUCCESS, - ip=ip, - started_at=utc_now, - finished_at=utc_now + timedelta(seconds=1), - raw_command=config.to_command_str(), - scan_group_id=scan_group_id, - config=config, - parsed=r, - ) - - -@pytest.fixture(scope="session") -def mtr_raw_output(request: Request) -> str: - fp = os.path.join(request.config.rootpath, "data/mtr_fatbeam.json") - with open(fp) as f: - data = f.read() - return data - - -@pytest.fixture(scope="session") -def mtr_result(mtr_raw_output: str) -> MTRResult: - return parse_mtr_output(mtr_raw_output, port=443, protocol=IPProtocol.TCP) - - -@pytest.fixture(scope="session") -def mtr_run(mtr_result: MTRResult, scan_group_id: str): - r = mtr_result - utc_now = datetime.now(tz=UTC) - config = MTRRunCommand( - command="mtr", - options=MTRRunCommandOptions( - ip=r.destination, protocol=IPProtocol.TCP, port=443 - ), - ) - - return MTRRun( - tool_version="1.2.3", - status=Status.SUCCESS, - ip=r.destination, - started_at=utc_now, - finished_at=utc_now + timedelta(seconds=1), - raw_command=config.to_command_str(), - scan_group_id=scan_group_id, - config=config, - parsed=r, - facility_id=1, - source_ip="1.2.3.4", - ) diff --git a/tests/managers/network/__init__.py b/tests/managers/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/managers/network/test_label.py b/tests/managers/network/test_label.py deleted file mode 100644 index abdd28f..0000000 --- a/tests/managers/network/test_label.py +++ /dev/null @@ -1,209 +0,0 @@ -import ipaddress -from datetime import datetime -from typing import TYPE_CHECKING - -import faker -import pytest -from psycopg.errors import UniqueViolation -from pydantic import ValidationError - -from generalresearch.models.network.label import ( - IPLabel, - IPLabelKind, - IPLabelMetadata, - IPLabelSource, -) -from generalresearch.models.thl.ipinfo import normalize_ip - -if TYPE_CHECKING: - from generalresearch.managers.network.label import IPLabelManager - -fake = faker.Faker() - - -@pytest.fixture -def ip_label(utc_now: datetime) -> IPLabel: - ip = ipaddress.IPv6Network((fake.ipv6(), 64), strict=False) - return IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip, - metadata=IPLabelMetadata(services=["RDP"]), - ) - - -def test_model(utc_now: datetime): - ip = fake.ipv4_public() - lbl = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip, - ) - assert lbl.ip.prefixlen == 32 - print(f"{lbl.ip=}") - - ip = ipaddress.IPv4Network((ip, 24), strict=False) - lbl = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip, - ) - print(f"{lbl.ip=}") - - with pytest.raises(ValidationError, match="IPv6 network must be /64 or larger"): - IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=fake.ipv6(), - ) - - ip = ipaddress.IPv6Network((fake.ipv6(), 64), strict=False) - lbl = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip, - ) - print(f"{lbl.ip=}") - - ip = ipaddress.IPv6Network((ip.network_address, 48), strict=False) - lbl = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip, - ) - print(f"{lbl.ip=}") - - -def test_create(iplabel_manager: IPLabelManager, ip_label: IPLabel): - iplabel_manager.create(ip_label) - - with pytest.raises( - UniqueViolation, match="duplicate key value violates unique constraint" - ): - iplabel_manager.create(ip_label) - - -def test_filter(iplabel_manager: IPLabelManager, ip_label: IPLabel, utc_hour_ago): - res = iplabel_manager.filter(ips=[ip_label.ip]) - assert len(res) == 0 - - iplabel_manager.create(ip_label) - res = iplabel_manager.filter(ips=[ip_label.ip]) - assert len(res) == 1 - - out = res[0] - assert out == ip_label - - res = iplabel_manager.filter(ips=[ip_label.ip], labeled_after=utc_hour_ago) - assert len(res) == 1 - - ip_label2 = ip_label.model_copy() - ip_label2.ip = fake.ipv4_public() - iplabel_manager.create(ip_label2) - res = iplabel_manager.filter(ips=[ip_label.ip, ip_label2.ip]) - assert len(res) == 2 - - -def test_filter_network( - iplabel_manager: IPLabelManager, ip_label: IPLabel, utc_hour_ago -): - print(ip_label) - ip_label = ip_label.model_copy() - ip_label.ip = ipaddress.IPv6Network((fake.ipv6(), 64), strict=False) - - iplabel_manager.create(ip_label) - res = iplabel_manager.filter(ips=[ip_label.ip]) - assert len(res) == 1 - - out = res[0] - assert out == ip_label - - res = iplabel_manager.filter(ips=[ip_label.ip], labeled_after=utc_hour_ago) - assert len(res) == 1 - - ip_label2 = ip_label.model_copy() - ip_label2.ip = fake.ipv4_public() - iplabel_manager.create(ip_label2) - res = iplabel_manager.filter(ips=[ip_label.ip, ip_label2.ip]) - assert len(res) == 2 - - -def test_network(iplabel_manager: IPLabelManager, utc_now: datetime): - # This is a fully-specific /128 ipv6 address. - # e.g. '51b7:b38d:8717:6c5b:cd3e:f5c3:3aba:17d' - ip = fake.ipv6() - # Generally, we'd want to annotate the /64 network - # e.g. '51b7:b38d:8717:6c5b::/64' - ip_64 = ipaddress.IPv6Network((ip, 64), strict=False) - - label = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip_64, - ) - iplabel_manager.create(label) - - # If I query for the /128 directly, I won't find it - res = iplabel_manager.filter(ips=[ip]) - assert len(res) == 0 - - # If I query for the /64 network I will - res = iplabel_manager.filter(ips=[ip_64]) - assert len(res) == 1 - - # Or, I can query for the /128 ip IN a network - res = iplabel_manager.filter(ip_in_network=ip) - assert len(res) == 1 - - -def test_label_cidr_and_ipinfo( - iplabel_manager: IPLabelManager, - ip_information_factory, - ip_geoname, - utc_now: datetime, -): - # We have network_iplabel.ip as a cidr col and - # thl_ipinformation.ip as a inet col. Make sure we can join appropriately - ip = fake.ipv6() - ip_information_factory(ip=ip, geoname=ip_geoname) - # We normalize for storage into ipinfo table - ip_norm, _ = normalize_ip(ip) - - # Test with a larger network - ip_48 = ipaddress.IPv6Network((ip, 48), strict=False) - print(f"{ip=}") - print(f"{ip_norm=}") - print(f"{ip_48=}") - label = IPLabel( - label_kind=IPLabelKind.VPN, - labeled_at=utc_now, - source=IPLabelSource.INTERNAL_USE, - provider="GeoNodE", - created_at=utc_now, - ip=ip_48, - ) - iplabel_manager.create(label) - - res = iplabel_manager.test_join(ip_norm) - print(res) diff --git a/tests/managers/network/test_tool_run.py b/tests/managers/network/test_tool_run.py deleted file mode 100644 index a815809..0000000 --- a/tests/managers/network/test_tool_run.py +++ /dev/null @@ -1,25 +0,0 @@ -def test_create_tool_run_from_nmap_run(nmap_run, toolrun_manager): - - toolrun_manager.create_nmap_run(nmap_run) - - run_out = toolrun_manager.get_nmap_run(nmap_run.id) - - assert nmap_run == run_out - - -def test_create_tool_run_from_rdns_run(rdns_run, toolrun_manager): - - toolrun_manager.create_rdns_run(rdns_run) - - run_out = toolrun_manager.get_rdns_run(rdns_run.id) - - assert rdns_run == run_out - - -def test_create_tool_run_from_mtr_run(mtr_run, toolrun_manager): - - toolrun_manager.create_mtr_run(mtr_run) - - run_out = toolrun_manager.get_mtr_run(mtr_run.id) - - assert mtr_run == run_out diff --git a/tests/models/network/__init__.py b/tests/models/network/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/models/network/test_mtr.py b/tests/models/network/test_mtr.py deleted file mode 100644 index 5d136c4..0000000 --- a/tests/models/network/test_mtr.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import faker - -from generalresearch.models.network.mtr.execute import execute_mtr -from generalresearch.models.network.tool_run import ToolClass, ToolName - -if TYPE_CHECKING: - from generalresearch.managers.network.tool_run import ToolRunManager - -fake = faker.Faker() - - -def test_execute_mtr(toolrun_manager: ToolRunManager): - ip = "65.19.129.53" - - run = execute_mtr(ip=ip, report_cycles=3) - assert run.tool_name == ToolName.MTR - assert run.tool_class == ToolClass.TRACEROUTE - assert run.ip == ip - result = run.parsed - - last_hop = result.hops[-1] - assert last_hop.asn == 6939 - assert last_hop.domain == "grlengine.com" - - last_hop_1 = result.hops[-2] - assert last_hop_1.asn == 6939 - assert last_hop_1.domain == "he.net" - - toolrun_manager.create_mtr_run(run) diff --git a/tests/models/network/test_nmap.py b/tests/models/network/test_nmap.py deleted file mode 100644 index 6adc9e4..0000000 --- a/tests/models/network/test_nmap.py +++ /dev/null @@ -1,39 +0,0 @@ -from __future__ import annotations - -import subprocess -from typing import TYPE_CHECKING - -import faker - -from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.nmap.execute import execute_nmap -from generalresearch.models.network.nmap.result import NmapResult, PortState -from generalresearch.models.network.tool_run import ToolClass, ToolName - -if TYPE_CHECKING: - from generalresearch.managers.network.tool_run import ToolRunManager - from generalresearch.models.network.tool_run import NmapRun - -fake = faker.Faker() - - -def resolve(host: str): - return subprocess.check_output(["dig", host, "+short"]).decode().strip() - - -def test_execute_nmap_scanme(toolrun_manager: ToolRunManager): - ip = resolve("scanme.nmap.org") - - run: NmapRun = execute_nmap( - ip=ip, top_ports=None, ports="20-30", enable_advanced=False - ) - assert run.tool_name == ToolName.NMAP - assert run.tool_class == ToolClass.PORT_SCAN - assert run.ip == ip - assert isinstance(run.parsed, NmapResult) - result = run.parsed - - port22 = result._port_index[(IPProtocol.TCP, 22)] - assert port22.state == PortState.OPEN - - toolrun_manager.create_nmap_run(run) diff --git a/tests/models/network/test_nmap_parser.py b/tests/models/network/test_nmap_parser.py deleted file mode 100644 index fc9884b..0000000 --- a/tests/models/network/test_nmap_parser.py +++ /dev/null @@ -1,32 +0,0 @@ -from __future__ import annotations - -import os -from typing import TYPE_CHECKING - -import pytest - -from generalresearch.models.network.nmap.parser import parse_nmap_xml -from generalresearch.models.network.nmap.result import NmapTrace - -if TYPE_CHECKING: - from generalresearch.models.network.nmap.result import NmapResult - - -@pytest.fixture -def nmap_raw_output_2(request) -> str: - fp = os.path.join(request.config.rootpath, "data/nmaprun2.xml") - with open(fp) as f: - data = f.read() - return data - - -def test_nmap_xml_parser(nmap_raw_output: str, nmap_raw_output_2: str): - n: NmapResult = parse_nmap_xml(nmap_raw_output) - assert n.tcp_open_ports == [61232] - - assert isinstance(n.trace, NmapTrace) - assert len(n.trace.hops) == 18 - - n = parse_nmap_xml(nmap_raw_output_2) - assert n.tcp_open_ports == [22, 80, 9929, 31337] - assert n.trace is None diff --git a/tests/models/network/test_rdns.py b/tests/models/network/test_rdns.py deleted file mode 100644 index 82126dd..0000000 --- a/tests/models/network/test_rdns.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import faker - -from generalresearch.models.network.rdns.execute import execute_rdns -from generalresearch.models.network.tool_run import ToolClass, ToolName - -if TYPE_CHECKING: - from generalresearch.managers.network.tool_run import ToolRunManager - -fake = faker.Faker() - - -def test_execute_rdns_grl(toolrun_manager: ToolRunManager): - ip = "65.19.129.53" - run = execute_rdns(ip=ip) - assert run.tool_name == ToolName.DIG - assert run.tool_class == ToolClass.RDNS - assert run.ip == ip - result = run.parsed - assert result.primary_hostname == "in1-smtp.grlengine.com" - assert result.primary_domain == "grlengine.com" - assert result.hostname_count == 1 - - toolrun_manager.create_rdns_run(run) - - -def test_execute_rdns_none(toolrun_manager: ToolRunManager): - ip = fake.ipv6() - run = execute_rdns(ip) - result = run.parsed - - assert result.primary_hostname is None - assert result.primary_domain is None - assert result.hostname_count == 0 - assert result.hostnames == [] - - toolrun_manager.create_rdns_run(run) -- cgit v1.2.3