diff options
| author | Max Nanis | 2026-08-17 02:54:18 -0700 |
|---|---|---|
| committer | Max Nanis | 2026-08-17 02:54:18 -0700 |
| commit | 7ff0d6fddb333dd973a827805c0c0cdcded220c7 (patch) | |
| tree | 528b2c61dedf96660d704a34a0081a1829dd474c | |
| parent | 7d466b47d732f16537370cb57ad52995aeef20f6 (diff) | |
| download | generalresearch-7ff0d6fddb333dd973a827805c0c0cdcded220c7.tar.gz generalresearch-7ff0d6fddb333dd973a827805c0c0cdcded220c7.zip | |
Ruff p2
97 files changed, 1306 insertions, 1299 deletions
diff --git a/generalresearch/models/legacy/api_status.py b/generalresearch/models/legacy/api_status.py index 8241eca..0422191 100644 --- a/generalresearch/models/legacy/api_status.py +++ b/generalresearch/models/legacy/api_status.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Literal, Optional +from typing import Literal from pydantic import BaseModel, Field @@ -14,7 +14,7 @@ class StatusResponse(BaseModel): status: Literal["success", "error"] = Field( description="The status of the API response.", examples=["success"] ) - msg: Optional[str] = Field( + msg: str | None = Field( description="An optional message, if the status is error.", examples=[""], default=None, diff --git a/generalresearch/models/legacy/bucket.py b/generalresearch/models/legacy/bucket.py index 8ce559b..52cbacf 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, Dict, List, Literal, Optional, Tuple, Union +from typing import Any, Literal from pydantic import ( BaseModel, @@ -57,9 +57,9 @@ class CategoryAssociation(BaseModel): examples=["People & Society"], ) - adwords_id: Optional[str] = Field(default=None, max_length=8, examples=["14"]) + adwords_id: str | None = Field(default=None, max_length=8, examples=["14"]) - adwords_label: Optional[str] = Field( + adwords_label: str | None = Field( default=None, max_length=255, examples=["People & Society"] ) @@ -151,7 +151,7 @@ class BucketBase(BaseModel): def censor(self): if not hasattr(self, "contents"): return - contents: List[BucketTask] = self.contents + contents: list[BucketTask] = self.contents for content in contents: content.censor() @@ -171,41 +171,41 @@ class Bucket(BaseModel): arbitrary_types_allowed=True, ) - name: Optional[str] = Field(default=None) - description: Optional[str] = Field(default=None) + name: str | None = Field(default=None) + description: str | None = Field(default=None) # pydantic serializes this to seconds - loi_min: Optional[timedelta] = Field(strict=True, default=None) - loi_max: Optional[timedelta] = Field(strict=True, default=None) - loi_mean: Optional[timedelta] = Field(strict=True, default=None) - loi_q1: Optional[timedelta] = Field(strict=True, default=None) - loi_q2: Optional[timedelta] = Field(strict=True, default=None) - loi_q3: Optional[timedelta] = Field(strict=True, default=None) + loi_min: timedelta | None = Field(strict=True, default=None) + loi_max: timedelta | None = Field(strict=True, default=None) + loi_mean: timedelta | None = Field(strict=True, default=None) + loi_q1: timedelta | None = Field(strict=True, default=None) + loi_q2: timedelta | None = Field(strict=True, default=None) + loi_q3: timedelta | None = Field(strict=True, default=None) # Decimal USD. This should not have more than 2 decimal places. # There is no way to make this "strict" and optional, so # we have a separate pre-validator - user_payout_min: Optional[Decimal] = Field(default=None, lt=1000, gt=0) - user_payout_max: Optional[Decimal] = Field(default=None, lt=1000, gt=0) - user_payout_q1: Optional[Decimal] = Field(default=None, lt=1000, gt=0) - user_payout_q2: Optional[Decimal] = Field(default=None, lt=1000, gt=0) - user_payout_q3: Optional[Decimal] = Field(default=None, lt=1000, gt=0) - user_payout_mean: Optional[Decimal] = Field(default=None, lt=1000, gt=0) + user_payout_min: Decimal | None = Field(default=None, lt=1000, gt=0) + user_payout_max: Decimal | None = Field(default=None, lt=1000, gt=0) + user_payout_q1: Decimal | None = Field(default=None, lt=1000, gt=0) + user_payout_q2: Decimal | None = Field(default=None, lt=1000, gt=0) + user_payout_q3: Decimal | None = Field(default=None, lt=1000, gt=0) + user_payout_mean: Decimal | None = Field(default=None, lt=1000, gt=0) - quality_score: Optional[float] = Field(default=None) + quality_score: float | None = Field(default=None) - category: List[CategoryAssociation] = Field(default_factory=list) + category: CategoryAssociation | None = Field(default_factory=list) - contents: Optional[List[BucketTask]] = Field(default=None) + contents: list[BucketTask] | None = Field(default=None) # This could store things like "is_recontact=False" - metadata: Dict[str, Union[str, float, bool, int]] = Field(default_factory=dict) + metadata: dict[str, str | float | bool | int] = Field(default_factory=dict) - eligibility_criteria: Optional[Tuple[SurveyEligibilityCriterion, ...]] = Field( + eligibility_criteria: tuple[SurveyEligibilityCriterion, ...] | None = Field( description="The reasons the user is eligible for tasks in this bucket", default=None, ) - eligibility_explanation: Optional[str] = Field( + eligibility_explanation: str | None = Field( default=None, description="Human-readable text explaining a user's eligibility for tasks in this bucket", examples=[ @@ -300,12 +300,12 @@ class Bucket(BaseModel): @field_validator("category") @classmethod - def check_category(cls, v: List[CategoryAssociation]) -> List[CategoryAssociation]: + def check_category(cls, v: list[CategoryAssociation]) -> list[CategoryAssociation]: assert sum(c.p for c in v) == 1, "sum of category score must be 1" return v @classmethod - def parse_from_offerwall(cls, bucket: Dict): + def parse_from_offerwall(cls, bucket: dict): """ This isn't really consistent across all offerwalls... Handle three cases: Could be {'payout': {'min': 123}}, or {'min_payout': 123} or {'payout': 123} @@ -324,7 +324,7 @@ class Bucket(BaseModel): return cls() @classmethod - def parse_from_offerwall_style1(cls, bucket: Dict): + def parse_from_offerwall_style1(cls, bucket: dict): # {'min_payout': 123} return cls( user_payout_min=cls.usd_cents_to_decimal(bucket["min_payout"]), @@ -360,7 +360,7 @@ class Bucket(BaseModel): ) @classmethod - def parse_from_offerwall_style2(cls, bucket: Dict[str, Any]): + def parse_from_offerwall_style2(cls, bucket: dict[str, Any]): # {'payout': {'min': 123}} loi_min_sec = bucket.get("duration", {}).get("min") loi_max_sec = bucket.get("duration", {}).get("max") @@ -385,7 +385,7 @@ class Bucket(BaseModel): ) @classmethod - def parse_from_offerwall_style3(cls, bucket: Dict[str, Any]): + def parse_from_offerwall_style3(cls, bucket: dict[str, Any]): # {'payout': 123, 'duration': 123} return cls( user_payout_min=cls.usd_cents_to_decimal(bucket["payout"]), @@ -399,13 +399,13 @@ class Bucket(BaseModel): ) @staticmethod - def usd_cents_to_decimal(v: Optional[int]) -> Optional[Decimal]: + def usd_cents_to_decimal(v: int | None) -> Decimal | None: if v is None: return None return Decimal(Decimal(int(v)) / Decimal(100)) @staticmethod - def decimal_to_usd_cents(d: Optional[Decimal]) -> Optional[Decimal]: + def decimal_to_usd_cents(d: Decimal | None) -> Decimal | None: if d is None: return None return round(d * Decimal(100), 2) @@ -421,7 +421,7 @@ class DurationSummary(StatisticalSummary): q1: int = Field(gt=0, le=60 * 90) q2: int = Field(gt=0, le=60 * 90) q3: int = Field(gt=0, le=60 * 90) - mean: Optional[int] = Field(gt=0, le=60 * 90, default=None) + mean: int | None = Field(gt=0, le=60 * 90, default=None) model_config = { "json_schema_extra": { @@ -460,7 +460,7 @@ class PayoutSummaryDecimal(StatisticalSummary): q1: Decimal = Field(gt=0, le=100) q2: Decimal = Field(gt=0, le=100) q3: Decimal = Field(gt=0, le=100) - mean: Optional[Decimal] = Field(gt=0, le=100, default=None) + mean: Decimal | None = Field(gt=0, le=100, default=None) class PayoutSummary(StatisticalSummary): @@ -471,7 +471,7 @@ class PayoutSummary(StatisticalSummary): q1: int = Field(gt=0, le=10000) q2: int = Field(gt=0, le=10000) q3: int = Field(gt=0, le=10000) - mean: Optional[int] = Field(gt=0, le=10000, default=None) + mean: int | None = Field(gt=0, le=10000, default=None) model_config = { "json_schema_extra": { @@ -514,45 +514,43 @@ class SurveyEligibilityCriterion(BaseModel): model_config = ConfigDict(validate_assignment=True) - question_id: Optional[UUIDStr] = Field( - examples=["71a367fb71b243dc89f0012e0ec91749"] - ) - property_code: Optional[PropertyCode] = Field(examples=["c:73629"]) + question_id: UUIDStr | None = Field(examples=["71a367fb71b243dc89f0012e0ec91749"]) + property_code: PropertyCode | None = Field(examples=["c:73629"]) question_text: str = Field( examples=[ "What company administers the retirement plan for your current employer?" ] ) # The answer(s) that were considered qualifying - qualifying_answer: Tuple[str, ...] = Field( + qualifying_answer: tuple[str, ...] = Field( description="User answer(s) that satisfied at least one eligibility rule", examples=["121"], ) - qualifying_answer_label: Optional[Tuple[str, ...]] = Field( + qualifying_answer_label: tuple[str, ...] | None = Field( examples=["Fidelity Investments"] ) - explanation: Optional[str] = Field( + explanation: str | None = Field( default=None, description="Human-readable text explaining how a user's answer to this question affects eligibility", examples=[ "The company that administers your employer's retirement plan is **Fidelity Investments**." ], ) - explanation_fragment: Optional[str] = Field( + explanation_fragment: str | None = Field( default=None, exclude=True, description="For internal use", examples=["who's retirement plan is administered by **Fidelity Investments**"], ) # Rank more "interesting"/rare/salient criterion first. - rank: Optional[NonNegativeInt] = Field( + rank: NonNegativeInt | None = Field( default=None, description="Lower values are shown more prominently in the UI", ) class TopNBucket(BucketBase): - category: List[CategoryAssociation] = Field(default_factory=list) + category: list[CategoryAssociation] = Field(default_factory=list) duration: DurationSummary = Field() payout: PayoutSummary = Field() quality_score: float = Field( @@ -586,8 +584,8 @@ class SingleEntryBucket(BucketBase): class TopNPlusBucket(BucketBase): - category: List[CategoryAssociation] = Field(default_factory=list) - contents: List[BucketTask] = Field() + category: list[CategoryAssociation] = Field(default_factory=list) + contents: list[BucketTask] = Field() duration: DurationSummary = Field() payout: PayoutSummary = Field() quality_score: float = Field() @@ -595,11 +593,11 @@ class TopNPlusBucket(BucketBase): description="This will always be 'USD'", default="USD", examples=["USD"] ) - eligibility_criteria: Tuple[SurveyEligibilityCriterion, ...] = Field( + eligibility_criteria: tuple[SurveyEligibilityCriterion, ...] = Field( description="The reasons the user is eligible for tasks in this bucket", default_factory=tuple, ) - eligibility_explanation: Optional[str] = Field( + eligibility_explanation: str | None = Field( default=None, description="Human-readable text explaining a user's eligibility for tasks in this bucket", examples=[ @@ -643,8 +641,8 @@ class TopNPlusBucket(BucketBase): class TopNPlusRecontactBucket(BucketBase): - category: List[CategoryAssociation] = Field(default_factory=list) - contents: List[BucketTask] = Field() + category: list[CategoryAssociation] = Field(default_factory=list) + contents: list[BucketTask] = Field() duration: DurationSummary = Field() payout: PayoutSummary = Field() quality_score: float = Field() @@ -670,17 +668,17 @@ class TopNPlusRecontactBucket(BucketBase): class SoftPairBucket(BucketBase): - uri: Optional[HttpsUrl] = Field( + uri: HttpsUrl | None = Field( examples=[None], description="The URL to send a respondent into. Must not edit this URL in any way. If the eligibility is " "conditional or ineligible, the uri will be null.", ) - category: List[CategoryAssociation] = Field(default_factory=list) - contents: List[BucketTask] = Field() + category: list[CategoryAssociation] = Field(default_factory=list) + contents: list[BucketTask] = Field() eligibility: Eligibility = Field(examples=["conditional"]) - missing_questions: List[str] = Field( + missing_questions: list[str] = Field( default_factory=list, examples=[["fb20fd4773304500b39c4f6de0012a5a"]] ) loi: int = Field(description="this is the max loi of the contents", examples=[612]) @@ -695,8 +693,8 @@ class SoftPairBucket(BucketBase): class MarketplaceBucket(BucketBase): - category: List[CategoryAssociation] = Field(default_factory=list) - contents: List[BucketTask] = Field() + category: list[CategoryAssociation] = Field(default_factory=list) + contents: list[BucketTask] = Field() duration: DurationSummary = Field() payout: PayoutSummary = Field() source: SourceName = Field( @@ -744,7 +742,7 @@ class OneShotOfferwallBucket(BaseModel): class OneShotSoftPairOfferwallBucket(OneShotOfferwallBucket): eligibility: Eligibility = Field(examples=["conditional"]) - missing_questions: List[str] = Field( + missing_questions: list[str] = Field( default_factory=list, examples=[["fb20fd4773304500b39c4f6de0012a5a"]] ) diff --git a/generalresearch/models/legacy/offerwall.py b/generalresearch/models/legacy/offerwall.py index a8efe9a..0e54387 100644 --- a/generalresearch/models/legacy/offerwall.py +++ b/generalresearch/models/legacy/offerwall.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Dict, List - from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt from generalresearch.models.custom_types import UUIDStr @@ -70,9 +68,9 @@ class OfferWall(BaseModel): default=0, ) - buckets: List[BucketBase] = Field(default_factory=list) + buckets: list[BucketBase] = Field(default_factory=list) - offerwall_reasons: List[OfferwallReason] = Field( + offerwall_reasons: list[OfferwallReason] = Field( default_factory=list, description=( "Explanations describing why so many or few opportunities are available." @@ -92,7 +90,7 @@ class SingleEntryOfferWall(OfferWall): """ payout_format: PayoutFormatType = PayoutFormatField - buckets: List[SingleEntryBucket] = Field(default_factory=list, max_length=1) + buckets: list[SingleEntryBucket] = Field(default_factory=list, max_length=1) class TopNOfferWall(OfferWall): @@ -102,7 +100,7 @@ class TopNOfferWall(OfferWall): Offerwall code: `45b7228a7` """ - buckets: List[TopNBucket] = Field(default_factory=list) + buckets: list[TopNBucket] = Field(default_factory=list) payout_format: PayoutFormatType = PayoutFormatField @@ -114,7 +112,7 @@ class StarwallOfferWall(OfferWall): Offerwall code: `b59a2d2b` """ - buckets: List[TopNBucket] = Field(default_factory=list) + buckets: list[TopNBucket] = Field(default_factory=list) payout_format: PayoutFormatType = PayoutFormatField @@ -124,7 +122,7 @@ class TopNPlusOfferWall(OfferWall): Offerwall code: `b145b803` """ - buckets: List[TopNPlusBucket] = Field(default_factory=list) + buckets: list[TopNPlusBucket] = Field(default_factory=list) class TopNPlusBlockOfferWall(OfferWall): @@ -134,7 +132,7 @@ class TopNPlusBlockOfferWall(OfferWall): Offerwall code: `d48cce47` """ - buckets: List[TopNPlusBucket] = Field(default_factory=list) + buckets: list[TopNPlusBucket] = Field(default_factory=list) # This incorrectly gets returned only when the user is blocked. It # shouldn't get returned at all @@ -149,7 +147,7 @@ class TopNPlusBlockRecontactOfferWall(OfferWall): Offerwall code: `1e5f0af8` """ - buckets: List[TopNPlusRecontactBucket] = Field(default_factory=list) + buckets: list[TopNPlusRecontactBucket] = Field(default_factory=list) # This incorrectly gets returned only when the user is blocked. It # shouldn't get returned at all @@ -162,7 +160,7 @@ class StarwallPlusOfferWall(OfferWall): Offerwall code: `5481f322` """ - buckets: List[TopNPlusBucket] = Field(default_factory=list) + buckets: list[TopNPlusBucket] = Field(default_factory=list) class StarwallPlusBlockOfferWall(OfferWall): @@ -172,7 +170,7 @@ class StarwallPlusBlockOfferWall(OfferWall): Offerwall code: `7fa1b3f4` """ - buckets: List[TopNPlusBucket] = Field(default_factory=list) + buckets: list[TopNPlusBucket] = Field(default_factory=list) # This incorrectly gets returned only when the user is blocked. It # shouldn't get returned at all @@ -187,7 +185,7 @@ class StarwallPlusBlockRecontactOfferWall(OfferWall): Offerwall code: `630db2a4` """ - buckets: List[TopNPlusRecontactBucket] = Field(default_factory=list) + buckets: list[TopNPlusRecontactBucket] = Field(default_factory=list) # This incorrectly gets returned only when the user is blocked. It # shouldn't get returned at all @@ -203,9 +201,9 @@ class SoftPairOfferwall(OfferWall): Offerwall code: `37d1da64` """ - buckets: List[SoftPairBucket] = Field(default_factory=list) + buckets: list[SoftPairBucket] = Field(default_factory=list) - question_info: Dict[str, "UpkQuestion"] = Field( + question_info: dict[str, "UpkQuestion"] = Field( default_factory=dict, examples=[ # { @@ -228,7 +226,7 @@ class MarketplaceOfferwall(OfferWall): Offerwall code: `5fa23085` """ - buckets: List[MarketplaceBucket] = Field(default_factory=list) + buckets: list[MarketplaceBucket] = Field(default_factory=list) class TimeBucksOfferwall(OfferWall): @@ -242,7 +240,7 @@ class TimeBucksOfferwall(OfferWall): Offerwall code: `1705e4f8` """ - buckets: List[TimeBucksBucket] = Field(default_factory=list) + buckets: list[TimeBucksBucket] = Field(default_factory=list) class TimeBucksBlockOfferwall(OfferWall): @@ -252,7 +250,7 @@ class TimeBucksBlockOfferwall(OfferWall): Offerwall code: `0af0f7ec` """ - buckets: List[TimeBucksBucket] = Field(default_factory=list) + buckets: list[TimeBucksBucket] = Field(default_factory=list) # This incorrectly gets returned only when the user is blocked. It shouldn't get returned at all payout_format: str = Field(exclude=True, default="") @@ -264,7 +262,7 @@ class OneShotOfferwall(OfferWall): Offerwall code: `6f27b1ae` """ - buckets: List[OneShotOfferwallBucket] = Field(default_factory=list) + buckets: list[OneShotOfferwallBucket] = Field(default_factory=list) class OneShotSoftPairOfferwall(SoftPairOfferwall): @@ -274,7 +272,7 @@ class OneShotSoftPairOfferwall(SoftPairOfferwall): Offerwall code: `18347426` """ - buckets: List[OneShotSoftPairOfferwallBucket] = Field(default_factory=list) + buckets: list[OneShotSoftPairOfferwallBucket] = Field(default_factory=list) class WXETOfferwall(OfferWall): @@ -282,7 +280,7 @@ class WXETOfferwall(OfferWall): Offerwall code: `55a4e1a9` """ - buckets: List[WXETOfferwallBucket] = Field(default_factory=list) + buckets: list[WXETOfferwallBucket] = Field(default_factory=list) class SingleEntryOfferWallResponse(OfferWallResponse): diff --git a/generalresearch/models/legacy/questions.py b/generalresearch/models/legacy/questions.py index 1a86121..04891ae 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, Dict, List, Optional +from typing import TYPE_CHECKING, Any from pydantic import ( BaseModel, @@ -32,11 +32,11 @@ if TYPE_CHECKING: class UpkQuestionResponse(StatusResponse): - questions: List[UpkQuestionOut] = Field() - consent_questions: List[Dict[str, Any]] = Field( + questions: list[UpkQuestionOut] = Field() + consent_questions: list[dict[str, Any]] = Field( description="For internal use", default_factory=list ) - special_questions: List[Dict[str, Any]] = Field( + special_questions: list[dict[str, Any]] = Field( description="For internal use", default_factory=list ) count: NonNegativeInt = Field(description="The number of questions returned") @@ -71,7 +71,7 @@ class UserQuestionAnswerIn(BaseModel): question_id: UUIDStr = Field(examples=["fb20fd4773304500b39c4f6de0012a5a"]) - answer: List[AnswerStr] = Field( + answer: list[AnswerStr] = Field( min_length=1, max_length=10, description="The user's answers to this question. Must pass the " @@ -113,7 +113,7 @@ class UserQuestionAnswerIn(BaseModel): @field_validator("answer", mode="after") @classmethod - def no_duplicate_answer_values(cls, v: List[AnswerStr]) -> List[AnswerStr]: + def no_duplicate_answer_values(cls, v: list[AnswerStr]) -> list[AnswerStr]: if len(v) != len(set(v)): raise ValueError("Don't provide duplicate answers") @@ -121,7 +121,7 @@ class UserQuestionAnswerIn(BaseModel): @field_validator("answer", mode="after") @classmethod - def sort_answer_values(cls, v: List[AnswerStr]) -> List[AnswerStr]: + def sort_answer_values(cls, v: list[AnswerStr]) -> list[AnswerStr]: return sorted(v) # --- Properties --- @@ -162,7 +162,7 @@ class UserQuestionAnswers(BaseModel): # POST /profiling-questions/ that they could use a randomly generated # session_id... I'm not sure, but it's entirely possible this will start # to cause issues in production. - session_id: Optional[UUIDStr] = Field( + session_id: UUIDStr | None = Field( default=None, description="The Session ID corresponds to the Wall.uuid. If profiling" "answers are being submitted directly, this can be None.", @@ -170,7 +170,7 @@ class UserQuestionAnswers(BaseModel): # We don't apply a default_factory here because there is no valid reason # why a GRS submission would come valid without any answers. - answers: Annotated[List[UserQuestionAnswerIn], BeforeValidator(preflight)] = Field( + answers: Annotated[list[UserQuestionAnswerIn], BeforeValidator(preflight)] = Field( min_length=1, max_length=100, description="The list of questions and their answers that are being" @@ -178,8 +178,8 @@ class UserQuestionAnswers(BaseModel): "(if via FSB).", ) - user: Optional[User] = Field(default=None) - wall: Optional[Wall] = Field(default=None) + user: User | None = Field(default=None) + wall: Wall | None = Field(default=None) # --- Validation --- @@ -209,7 +209,7 @@ class UserQuestionAnswers(BaseModel): @field_validator("answers", mode="after") @classmethod - def no_duplicate_questions(cls, v: List[UserQuestionAnswerIn]): + def no_duplicate_questions(cls, v: list[UserQuestionAnswerIn]): answer_qids = [qa.question_id for qa in v] if len(answer_qids) != len(set(answer_qids)): raise ValueError("Don't provide answers to duplicate questions") @@ -220,7 +220,7 @@ class UserQuestionAnswers(BaseModel): def prefetch_user(self, um: "UserManager") -> None: from generalresearch.models.thl.user import User - res: Optional[User] = um.get_user_if_exists( + res: User | None = um.get_user_if_exists( product_id=self.product_id, product_user_id=self.product_user_id ) @@ -228,13 +228,12 @@ class UserQuestionAnswers(BaseModel): raise ValidationError("Invalid user") self.user = res - return None def prefetch_wall(self, wm: "WallManager") -> None: from generalresearch.models import Source from generalresearch.models.thl.session import Wall - res: Optional[Wall] = wm.get_from_uuid_if_exists(wall_uuid=self.session_id) + res: Wall | None = wm.get_from_uuid_if_exists(wall_uuid=self.session_id) if res is None: raise ValueError("Invalid Event for session_id") @@ -252,4 +251,3 @@ class UserQuestionAnswers(BaseModel): raise ValueError("Not a valid GRS event status") self.wall = res - return None diff --git a/generalresearch/models/lucid/__init__.py b/generalresearch/models/lucid/__init__.py index 84a210d..c3365db 100644 --- a/generalresearch/models/lucid/__init__.py +++ b/generalresearch/models/lucid/__init__.py @@ -1,5 +1,4 @@ from pydantic import Field - from typing_extensions import Annotated LucidQuestionIdType = Annotated[ diff --git a/generalresearch/models/lucid/question.py b/generalresearch/models/lucid/question.py index 8542156..01c7859 100644 --- a/generalresearch/models/lucid/question.py +++ b/generalresearch/models/lucid/question.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional +from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, Field, field_validator, model_validator from typing_extensions import Self @@ -64,7 +64,7 @@ class LucidQuestion(MarketplaceQuestion): question_type: LucidQuestionType = Field( description="The type of question asked", frozen=True ) - options: Optional[List[LucidQuestionOption]] = Field( + options: list[LucidQuestionOption] | None = Field( default=None, min_length=1, frozen=True ) @@ -94,7 +94,7 @@ class LucidQuestion(MarketplaceQuestion): return options @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 = [ diff --git a/generalresearch/models/lucid/survey.py b/generalresearch/models/lucid/survey.py index 0e01243..4b1bb98 100644 --- a/generalresearch/models/lucid/survey.py +++ b/generalresearch/models/lucid/survey.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional, Self, Set, Tuple +from typing import Any, Self from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt @@ -23,7 +23,7 @@ class LucidCondition(MarketplaceCondition): id: BigAutoInteger = Field() source: Source = Field(default=Source.LUCID) - question_id: Optional[CoercedStr] = Field( + question_id: CoercedStr | None = Field( min_length=1, max_length=16, pattern=r"^[0-9]+$", @@ -41,7 +41,7 @@ class LucidCondition(MarketplaceCondition): return hash(self.id) @classmethod - def from_mysql(cls, x: Dict[str, Any]) -> Self: + def from_mysql(cls, x: dict[str, Any]) -> Self: x["value_type"] = ConditionValueType.LIST x["negate"] = False x["values"] = x.pop("pre_codes").split("|") @@ -58,11 +58,11 @@ class LucidQuota(BaseModel): id: BigAutoInteger = Field() uuid: UUIDStr = Field() upper_limit: NonNegativeInt = Field(examples=[20]) - criteria: List[int] = Field(min_length=1, max_length=25) + criteria: list[int] = Field(min_length=1, max_length=25) modified: AwareDatetimeISO = Field(description="modified or created") # We'll look this up with a special mysql query. If None, it means # that we don't know. - finish_count: Optional[int] = Field(default=None) + finish_count: int | None = Field(default=None) def __hash__(self): return hash(self.id) @@ -71,11 +71,11 @@ class LucidQuota(BaseModel): def is_open(self) -> bool: return self.upper_limit > self.finish_count - def passes(self, criteria_evaluation: Dict[int, Optional[bool]]) -> bool: + def passes(self, criteria_evaluation: dict[int, 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) - def matches(self, criteria_evaluation: Dict[int, Optional[bool]]) -> bool: + def matches(self, criteria_evaluation: dict[int, bool | None]) -> bool: # Matches means we meet all conditions. # We can "match" a quota that is closed. In that case, we would not be eligible for the survey. return all(criteria_evaluation.get(c) for c in self.criteria) @@ -93,8 +93,8 @@ class LucidQuota(BaseModel): # return True def matches_soft( - self, criteria_evaluation: Dict[int, Optional[bool]] - ) -> Tuple[Optional[bool], Set[int]]: + self, criteria_evaluation: dict[int, bool | None] + ) -> tuple[bool | None, set[int]]: # Passes back "matches" (T/F/none) and a list of unknown criterion hashes hash_evals = {cell: criteria_evaluation.get(cell) for cell in self.criteria} evals = set(hash_evals.values()) diff --git a/generalresearch/models/marketplace/summary.py b/generalresearch/models/marketplace/summary.py index 49cf7e9..f75c530 100644 --- a/generalresearch/models/marketplace/summary.py +++ b/generalresearch/models/marketplace/summary.py @@ -1,7 +1,8 @@ from __future__ import annotations from abc import ABC -from typing import Collection, Dict, List, Literal, Optional +from collections.abc import Collection +from typing import Literal import numpy as np from pydantic import BaseModel, ConfigDict, Field, computed_field @@ -16,7 +17,7 @@ class MarketplaceSummary(BaseModel): inventory: MarketplaceInventorySummary = Field( description="Inventory of the marketplace" ) - user_activity: Optional[str] = Field( + user_activity: str | None = Field( description="User activity of the marketplace", default=None ) @@ -24,32 +25,32 @@ class MarketplaceSummary(BaseModel): class MarketplaceInventorySummary(BaseModel): model_config = ConfigDict(validate_assignment=True) - live_tasks: List[CountStat] = Field( + live_tasks: list[CountStat] = Field( default_factory=list, description="The count of tasks that are currently live", ) - live_gen_pop_tasks: List[CountStat] = Field( + live_gen_pop_tasks: list[CountStat] = Field( default_factory=list, description="The count of gen-pop tasks that are currently live", ) - tasks_created: List[CountStat] = Field( + tasks_created: list[CountStat] = Field( default_factory=list, description="The count of tasks created", ) - required_finishes: List[CountStat] = Field( + required_finishes: list[CountStat] = Field( default_factory=list, description="Number of finishes needed across all live tasks", ) - payout: List[StatisticalSummaryStat] = Field( + payout: list[StatisticalSummaryStat] = Field( default_factory=list, description="The distribution of payouts for all live tasks", ) - expected_duration: List[StatisticalSummaryStat] = Field( + expected_duration: list[StatisticalSummaryStat] = Field( default_factory=list, description="The distribution of expected durations for all live tasks", ) - required_finishes_per_task: List[StatisticalSummaryStat] = Field( + required_finishes_per_task: list[StatisticalSummaryStat] = Field( default_factory=list, description="The distribution of required finishes on all live tasks", ) @@ -59,7 +60,7 @@ FacetKey = Literal["country_iso", "day", "month"] class Stat(BaseModel, ABC): - facet: Dict[FacetKey, str | int | float] = Field( + facet: dict[FacetKey, str | int | float] = Field( examples=[{"country_iso": "us"}], description="The grouping criteria" ) diff --git a/generalresearch/models/morning/question.py b/generalresearch/models/morning/question.py index 6dca8c0..7bc14ac 100644 --- a/generalresearch/models/morning/question.py +++ b/generalresearch/models/morning/question.py @@ -1,6 +1,6 @@ import json from enum import Enum -from typing import Any, Dict, List, Literal, Optional +from typing import Any, Literal from uuid import UUID from pydantic import BaseModel, Field, field_validator, model_validator diff --git a/generalresearch/models/morning/task_collection.py b/generalresearch/models/morning/task_collection.py index 8e1eb4b..b39f01b 100644 --- a/generalresearch/models/morning/task_collection.py +++ b/generalresearch/models/morning/task_collection.py @@ -1,4 +1,4 @@ -from typing import List, Set +from __future__ import annotations import pandas as pd from pandera import Check, Column, DataFrameSchema, Index @@ -11,8 +11,8 @@ from generalresearch.models.thl.survey.task_collection import ( create_empty_df_from_schema, ) -COUNTRY_ISOS: Set[str] = Localelator().get_all_countries() -LANGUAGE_ISOS: Set[str] = Localelator().get_all_languages() +COUNTRY_ISOS: set[str] = Localelator().get_all_countries() +LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() bid_stats_columns = { "system_conversion": Column(float, Check.between(0, 1), nullable=True), @@ -54,8 +54,8 @@ bid_columns = { } quota_columns = { "cpi": Column(float, Check.between(min_value=0, max_value=100)), - "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 } columns = ( @@ -84,7 +84,7 @@ MorningTaskCollectionSchema = DataFrameSchema( class MorningTaskCollection(TaskCollection): - items: List[MorningBid] + items: list[MorningBid] _schema = MorningTaskCollectionSchema def to_rows(self, bid: MorningBid): diff --git a/generalresearch/models/network/definitions.py b/generalresearch/models/network/definitions.py index 4fb44f4..2e1ab91 100644 --- a/generalresearch/models/network/definitions.py +++ b/generalresearch/models/network/definitions.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from enum import StrEnum from ipaddress import ip_address, ip_network -from typing import Optional CGNAT_NET = ip_network("100.64.0.0/10") @@ -36,7 +37,7 @@ class IPKind(StrEnum): UNSPECIFIED = "unspecified" -def get_ip_kind(ip: Optional[str]) -> Optional[IPKind]: +def get_ip_kind(ip: str | None) -> IPKind | None: if not ip: return None diff --git a/generalresearch/models/network/label.py b/generalresearch/models/network/label.py index b9a7659..e4ddd18 100644 --- a/generalresearch/models/network/label.py +++ b/generalresearch/models/network/label.py @@ -2,15 +2,14 @@ from __future__ import annotations import ipaddress from enum import StrEnum -from typing import Optional, List from pydantic import ( BaseModel, + ConfigDict, Field, + IPvAnyNetwork, computed_field, field_validator, - ConfigDict, - IPvAnyNetwork, ) from generalresearch.models.custom_types import ( @@ -69,7 +68,7 @@ class IPLabel(BaseModel): ip: IPvAnyNetwork = Field() labeled_at: AwareDatetimeISO = Field(default_factory=now_utc_factory) - created_at: Optional[AwareDatetimeISO] = Field(default=None) + created_at: AwareDatetimeISO | None = Field(default=None) label_kind: IPLabelKind = Field() source: IPLabelSource = Field() @@ -77,11 +76,11 @@ class IPLabel(BaseModel): 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: Optional[str] = Field( + provider: str | None = Field( default=None, examples=["geonode", "gecko"], max_length=128 ) - metadata: Optional[IPLabelMetadata] = Field(default=None) + metadata: IPLabelMetadata | None = Field(default=None) @field_validator("ip", mode="before") @classmethod @@ -96,7 +95,7 @@ class IPLabel(BaseModel): @field_validator("provider", mode="before") @classmethod - def provider_format(cls, v: Optional[str]) -> Optional[str]: + def provider_format(cls, v: str | None) -> str | None: if v is None: return v return v.lower().strip() @@ -123,4 +122,4 @@ class IPLabelMetadata(BaseModel): model_config = ConfigDict(validate_assignment=True, extra="allow") - services: Optional[List[str]] = Field(min_length=1, examples=[["RDP"]]) + services: list[str] | None = Field(min_length=1, examples=[["RDP"]]) diff --git a/generalresearch/models/network/mtr/command.py b/generalresearch/models/network/mtr/command.py index f8d2d49..fd5a8d0 100644 --- a/generalresearch/models/network/mtr/command.py +++ b/generalresearch/models/network/mtr/command.py @@ -1,5 +1,6 @@ +from __future__ import annotations + import subprocess -from typing import List, Optional from generalresearch.models.network.definitions import IPProtocol from generalresearch.models.network.mtr.parser import parse_mtr_output @@ -17,9 +18,9 @@ PROTOCOLS_W_PORT = {IPProtocol.TCP, IPProtocol.UDP, IPProtocol.SCTP} def build_mtr_command( ip: str, - protocol: Optional[IPProtocol] = None, - port: Optional[int] = None, - report_cycles: int = 10, + 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" @@ -67,4 +68,6 @@ def run_mtr(config: MTRRunCommand) -> MTRResult: check=False, ) raw = proc.stdout.strip() - return parse_mtr_output(raw, protocol=config.options.protocol, port=config.options.port) + 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 index a6fb82a..d77e814 100644 --- a/generalresearch/models/network/mtr/execute.py +++ b/generalresearch/models/network/mtr/execute.py @@ -1,15 +1,15 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Optional 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 ( - run_mtr, get_mtr_version, - build_mtr_command, + run_mtr, ) -from generalresearch.models.network.tool_run import MTRRun, ToolName, ToolClass, Status +from generalresearch.models.network.tool_run import MTRRun, Status, ToolClass, ToolName from generalresearch.models.network.tool_run_command import ( MTRRunCommand, MTRRunCommandOptions, @@ -19,9 +19,9 @@ from generalresearch.models.network.utils import get_source_ip def execute_mtr( ip: str, - scan_group_id: Optional[UUIDStr] = None, - protocol: Optional[IPProtocol] = IPProtocol.ICMP, - port: Optional[int] = None, + scan_group_id: UUIDStr | None = None, + protocol: IPProtocol | None = IPProtocol.ICMP, + port: int | None = None, report_cycles: int = 10, ) -> MTRRun: config = MTRRunCommand( diff --git a/generalresearch/models/network/mtr/parser.py b/generalresearch/models/network/mtr/parser.py index 685eeca..c29439e 100644 --- a/generalresearch/models/network/mtr/parser.py +++ b/generalresearch/models/network/mtr/parser.py @@ -1,5 +1,4 @@ import json -from typing import Dict from generalresearch.models.network.definitions import IPProtocol from generalresearch.models.network.mtr.result import MTRResult @@ -12,7 +11,7 @@ def parse_mtr_output(raw: str, port: int, protocol: IPProtocol) -> MTRResult: return MTRResult.model_validate(data) -def parse_mtr_raw_output(raw: str) -> Dict: +def parse_mtr_raw_output(raw: str) -> dict: data = json.loads(raw)["report"] data.update(data.pop("mtr")) data["hops"] = data.pop("hubs") diff --git a/generalresearch/models/network/mtr/result.py b/generalresearch/models/network/mtr/result.py index 5c775b4..34de845 100644 --- a/generalresearch/models/network/mtr/result.py +++ b/generalresearch/models/network/mtr/result.py @@ -1,19 +1,20 @@ +from __future__ import annotations + import re from functools import cached_property from ipaddress import ip_address -from typing import List, Optional import tldextract from pydantic import ( - Field, - field_validator, BaseModel, ConfigDict, - model_validator, + Field, computed_field, + field_validator, + model_validator, ) -from generalresearch.models.network.definitions import IPProtocol, get_ip_kind, IPKind +from generalresearch.models.network.definitions import IPKind, IPProtocol, get_ip_kind HOST_RE = re.compile(r"^(?P<hostname>.+?) \((?P<ip>[^)]+)\)$") @@ -23,7 +24,7 @@ class MTRHop(BaseModel): hop: int = Field(alias="count") host: str - asn: Optional[int] = Field(default=None, alias="ASN") + asn: int | None = Field(default=None, alias="ASN") loss_pct: float = Field(alias="Loss%") sent: int = Field(alias="Snt") @@ -34,10 +35,10 @@ class MTRHop(BaseModel): worst_ms: float = Field(alias="Wrst") stdev_ms: float = Field(alias="StDev") - hostname: Optional[str] = Field( + hostname: str | None = Field( default=None, examples=["fixed-187-191-8-145.totalplay.net"] ) - ip: Optional[str] = None + ip: str | None = None @field_validator("asn", mode="before") @classmethod @@ -74,7 +75,7 @@ class MTRHop(BaseModel): return self @cached_property - def ip_kind(self) -> Optional[IPKind]: + def ip_kind(self) -> IPKind | None: return get_ip_kind(self.ip) @cached_property @@ -85,7 +86,7 @@ class MTRHop(BaseModel): @computed_field(examples=["totalplay.net"]) @cached_property - def domain(self) -> Optional[str]: + def domain(self) -> str | None: if self.hostname: return tldextract.extract(self.hostname).top_domain_under_public_suffix @@ -120,9 +121,9 @@ class MTRResult(BaseModel): # Protocol used for the traceroute protocol: IPProtocol = Field(default=IPProtocol.ICMP) # The target port number for TCP/SCTP/UDP traces - port: Optional[int] = Field(default=None) + port: int | None = Field(default=None) - hops: List[MTRHop] = Field() + hops: list[MTRHop] = Field() def model_dump_postgres(self): # Writes for the network_mtr table diff --git a/generalresearch/models/network/nmap/command.py b/generalresearch/models/network/nmap/command.py index 47e0a87..6a524a8 100644 --- a/generalresearch/models/network/nmap/command.py +++ b/generalresearch/models/network/nmap/command.py @@ -1,5 +1,6 @@ +from __future__ import annotations + import subprocess -from typing import Optional, List from generalresearch.models.network.nmap.parser import parse_nmap_xml from generalresearch.models.network.nmap.result import NmapResult @@ -11,8 +12,8 @@ def build_nmap_command( no_ping: bool = True, enable_advanced: bool = True, timing: int = 4, - ports: Optional[str] = None, - top_ports: Optional[int] = None, + 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 diff --git a/generalresearch/models/network/nmap/execute.py b/generalresearch/models/network/nmap/execute.py index 0334f50..8a73307 100644 --- a/generalresearch/models/network/nmap/execute.py +++ b/generalresearch/models/network/nmap/execute.py @@ -1,9 +1,10 @@ -from typing import Optional +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, ToolName, ToolClass, Status +from generalresearch.models.network.tool_run import NmapRun, Status, ToolClass, ToolName from generalresearch.models.network.tool_run_command import ( NmapRunCommand, NmapRunCommandOptions, @@ -12,12 +13,12 @@ from generalresearch.models.network.tool_run_command import ( def execute_nmap( ip: str, - top_ports: Optional[int] = 1000, - ports: Optional[str] = None, + top_ports: int | None = 1000, + ports: str | None = None, no_ping: bool = True, enable_advanced: bool = True, timing: int = 4, - scan_group_id: Optional[UUIDStr] = None, + scan_group_id: UUIDStr | None = None, ): config = NmapRunCommand( options=NmapRunCommandOptions( diff --git a/generalresearch/models/network/nmap/parser.py b/generalresearch/models/network/nmap/parser.py index 5a441bb..49c13c7 100644 --- a/generalresearch/models/network/nmap/parser.py +++ b/generalresearch/models/network/nmap/parser.py @@ -1,26 +1,28 @@ +from __future__ import annotations + import xml.etree.cElementTree as ET from datetime import datetime, timezone -from typing import List, Dict, Any, Tuple, Optional +from typing import Any from generalresearch.models.network.definitions import IPProtocol from generalresearch.models.network.nmap.result import ( NmapHostname, - NmapResult, - NmapPort, - PortState, - PortStateReason, - NmapService, - NmapScript, - NmapPortStats, - NmapScanType, + NmapHostScript, NmapHostState, NmapHostStatusReason, - NmapHostScript, - NmapOSMatch, NmapOSClass, + NmapOSMatch, + NmapPort, + NmapPortStats, + NmapResult, + NmapScanInfo, + NmapScanType, + NmapScript, + NmapService, NmapTrace, NmapTraceHop, - NmapScanInfo, + PortState, + PortStateReason, ) @@ -109,7 +111,7 @@ class NmapXmlParser: return NmapScanInfo.model_validate(data) @classmethod - def _parse_runstats(cls, root: ET.Element) -> Dict: + def _parse_runstats(cls, root: ET.Element) -> dict: runstats = root.find("runstats") if runstats is None: return {} @@ -129,7 +131,7 @@ class NmapXmlParser: } @classmethod - def _parse_nmaprun(cls, nmaprun_el: ET.Element) -> Dict: + def _parse_nmaprun(cls, nmaprun_el: ET.Element) -> dict: nmap_data = dict() nmaprun = dict(nmaprun_el.attrib) nmap_data["command_line"] = nmaprun["args"] @@ -141,7 +143,7 @@ class NmapXmlParser: return nmap_data @classmethod - def _parse_xml_host(cls, host_el: ET.Element) -> Dict: + def _parse_xml_host(cls, host_el: ET.Element) -> dict: """ Receives a <host> XML tag representing a scanned host with its services. @@ -207,15 +209,15 @@ class NmapXmlParser: return data @classmethod - def _parse_os_matches(cls, host_el: ET.Element) -> List[NmapOSMatch] | None: + 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] = [] + matches: list[NmapOSMatch] = [] for m in os_elem.findall("osmatch"): - classes: List[NmapOSClass] = [] + classes: list[NmapOSClass] = [] for c in m.findall("osclass"): cpes = [e.text.strip() for e in c.findall("cpe") if e.text] @@ -245,7 +247,7 @@ class NmapXmlParser: return matches or None @classmethod - def _parse_hostnames(cls, hostnames_el: ET.Element) -> List[NmapHostname]: + def _parse_hostnames(cls, hostnames_el: ET.Element) -> list[NmapHostname]: """ Parses the hostnames element. e.g. <hostnames> @@ -269,11 +271,11 @@ class NmapXmlParser: @classmethod def _parse_xml_ports( cls, ports_elem: ET.Element - ) -> Tuple[List[NmapPort], NmapPortStats]: + ) -> tuple[list[NmapPort], NmapPortStats]: """ Parses the list of scanned services from a targeted host. """ - ports: List[NmapPort] = [] + ports: list[NmapPort] = [] stats = NmapPortStats() # handle extraports first @@ -319,7 +321,7 @@ class NmapXmlParser: "output": output, } - elements: Dict[str, Any] = {} + elements: dict[str, Any] = {} # handle <elem key="...">value</elem> for elem in script_elem.findall(".//elem"): @@ -373,7 +375,7 @@ class NmapXmlParser: return NmapPort.model_validate(port) @classmethod - def _parse_trace(cls, host_elem: ET.Element) -> Optional[NmapTrace]: + def _parse_trace(cls, host_elem: ET.Element) -> NmapTrace | None: trace_elem = host_elem.find("trace") if trace_elem is None: return None @@ -381,7 +383,7 @@ class NmapXmlParser: port_attr = trace_elem.attrib.get("port") proto_attr = trace_elem.attrib.get("proto") - hops: List[NmapTraceHop] = [] + hops: list[NmapTraceHop] = [] for hop_elem in trace_elem.findall("hop"): ttl = hop_elem.attrib.get("ttl") diff --git a/generalresearch/models/network/nmap/result.py b/generalresearch/models/network/nmap/result.py index d75efba..3f9cae6 100644 --- a/generalresearch/models/network/nmap/result.py +++ b/generalresearch/models/network/nmap/result.py @@ -1,10 +1,12 @@ +from __future__ import annotations + import json from datetime import timedelta from enum import StrEnum from functools import cached_property -from typing import Dict, Any, Literal, List, Optional, Tuple, Set +from typing import Any, Literal, Set -from pydantic import computed_field, BaseModel, Field +from pydantic import BaseModel, Field, computed_field from generalresearch.models.custom_types import AwareDatetimeISO, IPvAnyAddressStr from generalresearch.models.network.definitions import IPProtocol @@ -80,18 +82,18 @@ class NmapHostStatusReason(StrEnum): class NmapOSClass(BaseModel): vendor: str = None osfamily: str = None - osgen: Optional[str] = None + osgen: str | None = None accuracy: int = None - cpe: Optional[List[str]] = None + cpe: list[str] | None = None class NmapOSMatch(BaseModel): name: str accuracy: int - classes: List[NmapOSClass] = Field(default_factory=list) + classes: list[NmapOSClass] = Field(default_factory=list) @property - def best_class(self) -> Optional[NmapOSClass]: + def best_class(self) -> NmapOSClass | None: if not self.classes: return None return max(self.classes, key=lambda m: m.accuracy) @@ -108,19 +110,19 @@ class NmapScript(BaseModel): """ id: str - output: Optional[str] = None - elements: Dict[str, Any] = Field(default_factory=dict) + output: str | None = None + elements: dict[str, Any] = Field(default_factory=dict) class NmapService(BaseModel): # <service name="socks5" extrainfo="Username/password authentication required" method="probed" conf="10"/> - name: Optional[str] = None - product: Optional[str] = None - version: Optional[str] = None - extrainfo: Optional[str] = None - method: Optional[str] = None - conf: Optional[int] = None - cpe: List[str] = Field(default_factory=list) + 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") @@ -133,11 +135,11 @@ class NmapPort(BaseModel): protocol: IPProtocol = Field() # Closed ports will not have a NmapPort record state: PortState = Field() - reason: Optional[PortStateReason] = Field(default=None) - reason_ttl: Optional[int] = Field(default=None) + reason: PortStateReason | None = Field(default=None) + reason_ttl: int | None = Field(default=None) - service: Optional[NmapService] = None - scripts: List[NmapScript] = Field(default_factory=list) + 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 @@ -160,7 +162,7 @@ class NmapPort(BaseModel): class NmapHostScript(BaseModel): id: str = Field() - output: Optional[str] = Field(default=None) + output: str | None = Field(default=None) class NmapTraceHop(BaseModel): @@ -173,17 +175,17 @@ class NmapTraceHop(BaseModel): ttl: int = Field() - ipaddr: Optional[str] = Field( + ipaddr: str | None = Field( default=None, description="IP address of the responding router or host", ) - rtt_ms: Optional[float] = Field( + rtt_ms: float | None = Field( default=None, description="Round-trip time in milliseconds for the probe reaching this hop.", ) - host: Optional[str] = Field( + host: str | None = Field( default=None, description="Reverse DNS hostname for the hop if Nmap resolved one.", ) @@ -203,29 +205,29 @@ class NmapTrace(BaseModel): </trace> """ - port: Optional[int] = Field( + port: int | None = Field( default=None, description="Destination port used for traceroute probes (may be absent depending on scan type).", ) - protocol: Optional[IPProtocol] = Field( + protocol: IPProtocol | None = Field( default=None, description="Transport protocol used for the traceroute probes (tcp, udp, etc.).", ) - hops: List[NmapTraceHop] = Field( + hops: list[NmapTraceHop] = Field( default_factory=list, description="Ordered list of hops observed during the traceroute.", ) @property - def destination(self) -> Optional[NmapTraceHop]: + def destination(self) -> NmapTraceHop | None: return self.hops[-1] if self.hops else None class NmapHostname(BaseModel): # <hostname name="108-171-53-1.aceips.com" type="PTR"/> name: str - type: Optional[Literal["PTR", "user"]] = None + type: Literal["PTR", "user"] | None = None class NmapPortStats(BaseModel): @@ -280,11 +282,11 @@ class NmapResult(BaseModel): version: str = Field() xmloutputversion: str = Field() - scan_infos: List[NmapScanInfo] = Field(min_length=1) + scan_infos: list[NmapScanInfo] = Field(min_length=1) # comes from <runstats> - finished_at: Optional[AwareDatetimeISO] = Field(default=None) - exit_status: Optional[Literal["success", "error"]] = Field(default=None) + 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 @@ -293,53 +295,53 @@ class NmapResult(BaseModel): # <status state="up" reason="user-set" reason_ttl="0"/> host_state: NmapHostState = Field() host_state_reason: NmapHostStatusReason = Field() - host_state_reason_ttl: Optional[int] = None + host_state_reason_ttl: int | None = None # <address addr="108.171.53.1" addrtype="ipv4"/> target_ip: IPvAnyAddressStr = Field() - hostnames: List[NmapHostname] = Field() + hostnames: list[NmapHostname] = Field() - ports: List[NmapPort] = [] + ports: list[NmapPort] = [] port_stats: NmapPortStats = Field() # <uptime seconds="4063775" lastboot="Fri Jan 16 12:12:06 2026"/> - uptime_seconds: Optional[int] = Field(default=None) + uptime_seconds: int | None = Field(default=None) # <distance value="11"/> - distance: Optional[int] = Field(description="approx number of hops", default=None) + distance: int | None = Field(description="approx number of hops", default=None) # <tcpsequence index="263" difficulty="Good luck!"> - tcp_sequence_index: Optional[int] = None - tcp_sequence_difficulty: Optional[str] = None + tcp_sequence_index: int | None = None + tcp_sequence_difficulty: str | None = None # <ipidsequence class="All zeros"> - ipid_sequence_class: Optional[str] = None + ipid_sequence_class: str | None = None # <tcptssequence class="1000HZ" > - tcp_timestamp_class: Optional[str] = None + tcp_timestamp_class: str | None = None # <times srtt="54719" rttvar="23423" to="148411"/> - srtt_us: Optional[int] = Field( + srtt_us: int | None = Field( default=None, description="smoothed RTT estimate (microseconds µs)" ) - rttvar_us: Optional[int] = Field( + rttvar_us: int | None = Field( default=None, description="RTT variance (microseconds µs)" ) - timeout_us: Optional[int] = Field( + timeout_us: int | None = Field( default=None, description="probe timeout (microseconds µs)" ) - os_matches: Optional[List[NmapOSMatch]] = Field(default=None) + os_matches: list[NmapOSMatch] | None = Field(default=None) - host_scripts: List[NmapHostScript] = Field(default_factory=list) + host_scripts: list[NmapHostScript] = Field(default_factory=list) - trace: Optional[NmapTrace] = Field(default=None) + trace: NmapTrace | None = Field(default=None) - raw_xml: Optional[str] = None + raw_xml: str | None = None @computed_field @property - def last_boot(self) -> Optional[AwareDatetimeISO]: + def last_boot(self) -> AwareDatetimeISO | None: if self.uptime_seconds: return self.started_at - timedelta(seconds=self.uptime_seconds) @@ -356,20 +358,20 @@ class NmapResult(BaseModel): ) @property - def latency_ms(self) -> Optional[float]: + 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) -> Optional[NmapOSMatch]: + 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]: + 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]: + def tcp_open_ports(self) -> list[int]: """ Returns a list of open TCP port numbers. """ @@ -379,7 +381,7 @@ class NmapResult(BaseModel): ] @property - def udp_open_ports(self) -> List[int]: + def udp_open_ports(self) -> list[int]: """ Returns a list of open UDP port numbers. """ @@ -389,7 +391,7 @@ class NmapResult(BaseModel): ] @cached_property - def _port_index(self) -> Dict[Tuple[IPProtocol, int], NmapPort]: + def _port_index(self) -> dict[tuple[IPProtocol, int], NmapPort]: return {(p.protocol, p.port): p for p in self.ports} def get_port_state( diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py index 03a5080..cabd13c 100644 --- a/generalresearch/models/network/rdns/execute.py +++ b/generalresearch/models/network/rdns/execute.py @@ -1,18 +1,18 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Optional from uuid import uuid4 from generalresearch.models.custom_types import UUIDStr from generalresearch.models.network.rdns.command import ( - run_rdns, get_dig_version, - build_rdns_command, + run_rdns, ) from generalresearch.models.network.tool_run import ( - ToolName, - ToolClass, - Status, RDNSRun, + Status, + ToolClass, + ToolName, ) from generalresearch.models.network.tool_run_command import ( RDNSRunCommand, @@ -20,7 +20,7 @@ from generalresearch.models.network.tool_run_command import ( ) -def execute_rdns(ip: str, scan_group_id: Optional[UUIDStr] = None): +def execute_rdns(ip: str, scan_group_id: UUIDStr | None = None): started_at = datetime.now(tz=timezone.utc) tool_version = get_dig_version() config = RDNSRunCommand(options=RDNSRunCommandOptions(ip=ip)) diff --git a/generalresearch/models/network/rdns/parser.py b/generalresearch/models/network/rdns/parser.py index 231949e..e1cf023 100644 --- a/generalresearch/models/network/rdns/parser.py +++ b/generalresearch/models/network/rdns/parser.py @@ -1,6 +1,5 @@ import ipaddress import re -from typing import List from generalresearch.models.custom_types import IPvAnyAddressStr from generalresearch.models.network.rdns.result import RDNSResult @@ -8,8 +7,8 @@ from generalresearch.models.network.rdns.result import RDNSResult PTR_RE = re.compile(r"\sPTR\s+([^\s]+)\.") -def parse_rdns_output(ip: IPvAnyAddressStr, raw:str): - hostnames: List[str] = [] +def parse_rdns_output(ip: IPvAnyAddressStr, raw: str): + hostnames: list[str] = [] for line in raw.splitlines(): m = PTR_RE.search(line) diff --git a/generalresearch/models/network/rdns/result.py b/generalresearch/models/network/rdns/result.py index 81b4085..46af643 100644 --- a/generalresearch/models/network/rdns/result.py +++ b/generalresearch/models/network/rdns/result.py @@ -1,9 +1,10 @@ +from __future__ import annotations + import json from functools import cached_property -from typing import Optional, List import tldextract -from pydantic import BaseModel, Field, model_validator, computed_field +from pydantic import BaseModel, Field, computed_field, model_validator from generalresearch.models.custom_types import IPvAnyAddressStr @@ -12,7 +13,7 @@ class RDNSResult(BaseModel): ip: IPvAnyAddressStr = Field() - hostnames: List[str] = Field(default_factory=list) + hostnames: list[str] = Field(default_factory=list) @model_validator(mode="after") def validate_hostname_prop(self): @@ -24,7 +25,7 @@ class RDNSResult(BaseModel): @computed_field(examples=["fixed-187-191-8-145.totalplay.net"]) @cached_property - def primary_hostname(self) -> Optional[str]: + def primary_hostname(self) -> str | None: if self.hostnames: return self.hostnames[0] @@ -35,7 +36,7 @@ class RDNSResult(BaseModel): @computed_field(examples=["totalplay.net"]) @cached_property - def primary_domain(self) -> Optional[str]: + def primary_domain(self) -> str | None: if self.primary_hostname: return tldextract.extract( self.primary_hostname diff --git a/generalresearch/models/network/tool_run.py b/generalresearch/models/network/tool_run.py index 114d4b6..c49ffc0 100644 --- a/generalresearch/models/network/tool_run.py +++ b/generalresearch/models/network/tool_run.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from enum import StrEnum -from typing import Optional, Literal +from typing import Literal from uuid import uuid4 from pydantic import BaseModel, Field, PositiveInt @@ -13,10 +15,10 @@ 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 ( - ToolRunCommand, + MTRRunCommand, NmapRunCommand, RDNSRunCommand, - MTRRunCommand, + ToolRunCommand, ) @@ -48,7 +50,7 @@ class ToolRun(BaseModel): A run of a networking tool against one host/ip. """ - id: Optional[PositiveInt] = Field(default=None) + id: PositiveInt | None = Field(default=None) ip: IPvAnyAddressStr = Field() scan_group_id: UUIDStr = Field(default_factory=lambda: uuid4().hex) @@ -57,8 +59,8 @@ class ToolRun(BaseModel): tool_version: str = Field() started_at: AwareDatetimeISO = Field() - finished_at: Optional[AwareDatetimeISO] = Field(default=None) - status: Optional[Status] = Field(default=None) + finished_at: AwareDatetimeISO | None = Field(default=None) + status: Status | None = Field(default=None) raw_command: str = Field() diff --git a/generalresearch/models/network/tool_run_command.py b/generalresearch/models/network/tool_run_command.py index 68d2070..6f22d6b 100644 --- a/generalresearch/models/network/tool_run_command.py +++ b/generalresearch/models/network/tool_run_command.py @@ -1,4 +1,6 @@ -from typing import Dict, Optional, Literal +from __future__ import annotations + +from typing import Literal from pydantic import BaseModel, Field @@ -8,13 +10,13 @@ from generalresearch.models.network.definitions import IPProtocol class ToolRunCommand(BaseModel): command: str = Field() - options: Dict[str, Optional[str | int]] = Field(default_factory=dict) + options: dict[str, str | int | None] = Field(default_factory=dict) class NmapRunCommandOptions(BaseModel): ip: IPvAnyAddressStr - top_ports: Optional[int] = Field(default=1000) - ports: Optional[str] = Field(default=None) + 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) @@ -49,7 +51,7 @@ class RDNSRunCommand(ToolRunCommand): class MTRRunCommandOptions(BaseModel): ip: IPvAnyAddressStr = Field() protocol: IPProtocol = Field(default=IPProtocol.ICMP) - port: Optional[int] = Field(default=None) + port: int | None = Field(default=None) report_cycles: int = Field(default=10) diff --git a/generalresearch/models/pollfish/question.py b/generalresearch/models/pollfish/question.py index a89a793..9508154 100644 --- a/generalresearch/models/pollfish/question.py +++ b/generalresearch/models/pollfish/question.py @@ -1,8 +1,10 @@ +from __future__ import annotations + # https://wss.pollfish.com/mediation/documentation import json import logging from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Self +from typing import TYPE_CHECKING, Any, Literal, Self from pydantic import BaseModel, Field, model_validator @@ -59,9 +61,9 @@ class PollfishQuestion(MarketplaceQuestion): max_length=1024, min_length=1, description="The text shown to respondents" ) question_type: PollfishQuestionType = Field(frozen=True) - options: Optional[List[PollfishQuestionOption]] = Field(default=None, min_length=1) + options: list[PollfishQuestionOption] | None = Field(default=None, min_length=1) # This comes from the API field "category" - tags: Optional[str] = Field(default=None, frozen=True) + tags: str | None = Field(default=None, frozen=True) source: Literal[Source.POLLFISH] = Source.POLLFISH @property @@ -78,7 +80,7 @@ class PollfishQuestion(MarketplaceQuestion): return self @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 = [ @@ -97,7 +99,7 @@ class PollfishQuestion(MarketplaceQuestion): category_id=d.get("category_id"), ) - 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/precision/definitions.py b/generalresearch/models/precision/definitions.py index daa6f64..bd6a9c7 100644 --- a/generalresearch/models/precision/definitions.py +++ b/generalresearch/models/precision/definitions.py @@ -1,7 +1,8 @@ +from __future__ import annotations + # These were sent to us in an excel file. Copied here because we don't use excel files # also available here: https://integrations.precisionsample.com/api.html#API%20Lookup%20Document # https://integrations.precisionsample.com/PS_GetProjects_API_Lookup_values.xlsx - from generalresearch.locales import Localelator locales = Localelator() diff --git a/generalresearch/models/precision/question.py b/generalresearch/models/precision/question.py index 9673f54..d508086 100644 --- a/generalresearch/models/precision/question.py +++ b/generalresearch/models/precision/question.py @@ -1,8 +1,10 @@ +from __future__ import annotations + # https://integrations.precisionsample.com/api.html#Get%20Questions import json import logging from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Self +from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, Field, field_validator, model_validator @@ -65,7 +67,7 @@ class PrecisionQuestionType(str, Enum): class PrecisionUserQuestionAnswer(MarketplaceUserQuestionAnswer): question_id: PrecisionQuestionID = Field() - question_type: Optional[PrecisionQuestionType] = Field(default=None) + question_type: PrecisionQuestionType | None = Field(default=None) # Was this answer synchronized with precision's user profile API? synced: bool = Field(default=False) @@ -74,16 +76,16 @@ class PrecisionQuestion(MarketplaceQuestion): question_id: PrecisionQuestionID = Field( description="The unique identifier for the qualification" ) - question_name: Optional[str] = Field(default=None, max_length=128) + question_name: str | None = Field(default=None, max_length=128) question_text: str = Field( max_length=1024, min_length=1, description="The text shown to respondents" ) question_type: PrecisionQuestionType = Field(frozen=True) - options: Optional[List[PrecisionQuestionOption]] = Field(default=None, min_length=1) + options: list[PrecisionQuestionOption] | None = Field(default=None, min_length=1) # This comes from the API field ProfileName. idk what the possible values are, looks like: # 'Personal Profile', 'Work Profile', 'Auto Profile', 'Medical Profile', 'Travel & Entertainment'. # I don't know what, if anything, this is used for. - profile: Optional[str] = Field(default=None, frozen=True) + profile: str | None = Field(default=None, frozen=True) source: Literal[Source.PRECISION] = Source.PRECISION @property @@ -91,7 +93,7 @@ class PrecisionQuestion(MarketplaceQuestion): return self.question_id @field_validator("question_text", mode="after") - def remove_nbsp(cls, s: Optional[str]): + def remove_nbsp(cls, s: str | None): return string_utils.remove_nbsp(s) @model_validator(mode="after") @@ -104,7 +106,7 @@ class PrecisionQuestion(MarketplaceQuestion): return self @classmethod - def from_api(cls, d: Dict[str, Any]) -> Optional["PrecisionQuestion"]: + def from_api(cls, d: dict[str, Any]) -> "PrecisionQuestion" | None: """ :param d: Raw response from API """ @@ -115,7 +117,7 @@ class PrecisionQuestion(MarketplaceQuestion): return None @classmethod - def _from_api(cls, d: Dict[str, Any]) -> "PrecisionQuestion": + def _from_api(cls, d: dict[str, Any]) -> "PrecisionQuestion": question_type = PrecisionQuestionType.from_api(d["question_type_name"]) # sometimes an empty option is returned .... ? options = [ @@ -137,7 +139,7 @@ class PrecisionQuestion(MarketplaceQuestion): ) @classmethod - def from_db(cls, d: Dict[str, Any]) -> "PrecisionQuestion": + def from_db(cls, d: dict[str, Any]) -> "PrecisionQuestion": options = None if d["options"]: options = [ @@ -156,7 +158,7 @@ class PrecisionQuestion(MarketplaceQuestion): category_id=d.get("category_id"), ) - 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/prodege/question.py b/generalresearch/models/prodege/question.py index e3cdf08..f8741c9 100644 --- a/generalresearch/models/prodege/question.py +++ b/generalresearch/models/prodege/question.py @@ -6,7 +6,7 @@ import logging from datetime import datetime, timezone from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Self, Set +from typing import TYPE_CHECKING, Any, Literal from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator @@ -33,13 +33,13 @@ class ProdegeUserQuestionAnswer(BaseModel): # for "anonymous" users, which are represented by a list of question # answers not associated with an actual user. No default b/c we must # explicitly set the field to None. - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32) + user_id: PositiveInt | None = Field(lt=MAX_INT32) question_id: ProdegeQuestionIdType = Field() # This is optional b/c we do not need it when writing these to the # db. When these are fetched from the db for use in yield-management, # we read this field from the prodege_question table. - question_type: Optional[ProdegeQuestionType] = Field(default=None) + question_type: ProdegeQuestionType | None = Field(default=None) # 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]*$") @@ -58,10 +58,10 @@ class ProdegeUserQuestionAnswer(BaseModel): ) @cached_property - def options_ids(self) -> Set[str]: + def options_ids(self) -> set[str]: return set(self.option_id.split("|")) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"question_type"}) d["created"] = self.created.replace(tzinfo=None) return d @@ -122,8 +122,8 @@ class ProdegeQuestion(MarketplaceQuestion): question_text: str = Field(max_length=1024, min_length=1) question_type: ProdegeQuestionType = Field(frozen=True) # This comes from the API category, but is not great (most are "Consumer Lifestyle") - tags: Optional[str] = Field(default=None, frozen=True) - options: Optional[List[ProdegeQuestionOption]] = Field(default=None, min_length=1) + tags: str | None = Field(default=None, frozen=True) + options: list[ProdegeQuestionOption] | None = Field(default=None, min_length=1) source: Literal[Source.PRODEGE] = Source.PRODEGE @property @@ -139,9 +139,7 @@ class ProdegeQuestion(MarketplaceQuestion): return self @classmethod - def from_api( - cls, d: Dict[str, Any], country_iso: str - ) -> Optional["ProdegeQuestion"]: + def from_api(cls, d: dict[str, Any], country_iso: str) -> "ProdegeQuestion" | None: """ :param d: Raw response from API """ @@ -152,7 +150,7 @@ class ProdegeQuestion(MarketplaceQuestion): return None @classmethod - def _from_api(cls, d: Dict[str, Any], country_iso: str) -> "ProdegeQuestion": + def _from_api(cls, d: dict[str, Any], country_iso: str) -> "ProdegeQuestion": # The API has no concept of language at all. Questions for a country # are returned both in english and other languages. Questions do have # a field 'country_specific', and if True, that generally means the @@ -185,7 +183,7 @@ class ProdegeQuestion(MarketplaceQuestion): return cls.model_validate(d) @classmethod - def from_db(cls, d: Dict[str, Any]) -> "ProdegeQuestion": + def from_db(cls, d: dict[str, Any]) -> "ProdegeQuestion": options = None if d["options"]: options = [ @@ -210,7 +208,7 @@ class ProdegeQuestion(MarketplaceQuestion): tags=d.get("tags"), ) - 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/prodege/survey.py b/generalresearch/models/prodege/survey.py index eaae883..76af35d 100644 --- a/generalresearch/models/prodege/survey.py +++ b/generalresearch/models/prodege/survey.py @@ -7,7 +7,7 @@ from collections import defaultdict from datetime import datetime, timezone from decimal import Decimal from functools import cached_property -from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Type +from typing import Any, Literal, Type from pydantic import ( BaseModel, @@ -52,10 +52,10 @@ class ProdegeCondition(MarketplaceCondition): model_config = ConfigDict(populate_by_name=True) question_id: ProdegeQuestionIdType = Field() - values: List[str] = Field(validation_alias="precodes") + values: list[str] = Field(validation_alias="precodes") @classmethod - def from_api(cls, d: Dict[str, Any]) -> "ProdegeCondition": + def from_api(cls, d: dict[str, Any]) -> "ProdegeCondition": assert d["operator"] in { "OR", "NOT", @@ -97,7 +97,7 @@ class ProdegeQuota(BaseModel): description="The total number of allowed responses that remain from the sample_size", validation_alias="number_of_respondents", ) - condition_hashes: List[str] = Field(min_length=0, default_factory=list) + condition_hashes: list[str] = Field(min_length=0, default_factory=list) # Each quota can have a different calculation type, instead of on the survey calculation_type: TaskCalculationType = Field( description="Indicates whether the targets are counted per Complete or Survey Start", @@ -105,10 +105,10 @@ class ProdegeQuota(BaseModel): ) quota_id: CoercedStr = Field() # If the parent_quota_id is None, then this is a parent. There can be multiple parent quotas. - parent_quota_id: Optional[CoercedStr] = Field() + parent_quota_id: CoercedStr | None = Field() # ISO 3166-1 alpha-2 (two-letter codes, lowercase) - country_iso: Optional[str] = Field( + country_iso: str | None = Field( max_length=2, min_length=2, pattern=r"^[a-z]{2}$", default=None ) @@ -135,7 +135,7 @@ class ProdegeQuota(BaseModel): return "1" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: ProdegeCondition( question_id="3", @@ -151,7 +151,7 @@ class ProdegeQuota(BaseModel): } @classmethod - def from_api(cls, d: Dict[str, Any]) -> "ProdegeQuota": + def from_api(cls, d: dict[str, Any]) -> "ProdegeQuota": # the API doesn't handle None's correctly? idk if d["parent_quota_id"] == 0: d["parent_quota_id"] = None @@ -166,7 +166,7 @@ class ProdegeQuota(BaseModel): return cls.model_validate(d) def passes( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str + self, criteria_evaluation: dict[str, bool | None], country_iso: str ) -> bool: # Passes means we 1) meet all conditions (aka "match") AND 2) the quota is open. return self.is_open and self.matches( @@ -174,7 +174,7 @@ class ProdegeQuota(BaseModel): ) def matches( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str + self, criteria_evaluation: dict[str, bool | None], country_iso: str ) -> bool: # Match means we meet all conditions. # We can "match" a quota that is closed. In that case, we would @@ -187,7 +187,7 @@ class ProdegeQuota(BaseModel): return self.country_iso is None or self.country_iso == country_iso def passes_verbose( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str + self, criteria_evaluation: dict[str, bool | None], country_iso: str ) -> bool: print(f"quota.is_open: {self.is_open}") print( @@ -202,8 +202,8 @@ class ProdegeQuota(BaseModel): ) def passes_soft( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str - ) -> Tuple[Optional[bool], Set[str]]: + self, criteria_evaluation: dict[str, bool | None], country_iso: str + ) -> tuple[bool | None, set[str]]: # Passes back "passes" (T/F/none) and a list of unknown criterion hashes if self.is_open is False: return False, set() @@ -246,10 +246,10 @@ class ProdegeUserPastParticipation(BaseModel): survey_id: str = Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") started: AwareDatetimeISO = Field() # This is what is returned in the redirect in the url param "status". - ext_status_code_1: Optional[ProdgeRedirectStatus] = Field(default=None) + ext_status_code_1: ProdgeRedirectStatus | None = Field(default=None) @property - def participation_types(self) -> Set[ProdegePastParticipationType]: + def participation_types(self) -> set[ProdegePastParticipationType]: # If the survey is filtering completes, then only a complete # counts. But if the survey is filtering on clicks, then a person # who got a complete ALSO did click. And so, the logic here is that @@ -286,8 +286,8 @@ class ProdegePastParticipation(BaseModel): survey_ids: AlphaNumStrSet = Field(validation_alias="participation_project_ids") filter_type: InclExcl = Field() # API has a mistake. We treat 0 as null - in_past_days: Optional[int] = Field(default=None) - participation_types: List[ProdegePastParticipationType] = Field() + in_past_days: int | None = Field(default=None) + participation_types: list[ProdegePastParticipationType] = Field() """ e.g. Anyone who got a complete in either of these projects in the past 7 days, @@ -299,7 +299,7 @@ class ProdegePastParticipation(BaseModel): """ @classmethod - def from_api(cls, d: Dict[str, Any]) -> "ProdegePastParticipation": + def from_api(cls, d: dict[str, Any]) -> "ProdegePastParticipation": # the API doesn't handle None's correctly? idk if d["in_past_days"] == 0: d["in_past_days"] = None @@ -323,7 +323,7 @@ class ProdegePastParticipation(BaseModel): ) def is_eligible( - self, user_participations: List[ProdegeUserPastParticipation] + self, user_participations: list[ProdegeUserPastParticipation] ) -> bool: if self.filter_type == "include": # User is only eligible if they HAVE participated. Return True as soon as they match anything. @@ -383,12 +383,12 @@ class ProdegeSurvey(MarketplaceTask): # have to store it in the db if we see it. In API res, these are called # "loi" and "actual_ir", but the actual IR is only actually the actual # IR if the "phases" is "actual" :facepalm: - bid_loi: Optional[int] = Field(default=None, le=120 * 60) - bid_ir: Optional[float] = Field(default=None, ge=0, le=1) - actual_loi: Optional[int] = Field(default=None, le=120 * 60) - actual_ir: Optional[float] = Field(default=None, ge=0, le=1) + bid_loi: int | None = Field(default=None, le=120 * 60) + bid_ir: float | None = Field(default=None, ge=0, le=1) + actual_loi: int | None = Field(default=None, le=120 * 60) + actual_ir: float | None = Field(default=None, ge=0, le=1) # Unclear what the difference is bw IR and conversion - conversion_rate: Optional[float] = Field(default=None, ge=0, le=1) + conversion_rate: float | None = Field(default=None, ge=0, le=1) entrance_url: str = Field( description="The link survey respondents should be sent to", @@ -396,12 +396,12 @@ class ProdegeSurvey(MarketplaceTask): ) # This described time-based click rate limiting. - max_clicks_settings: Optional[ProdegeMaxClicksSetting] = Field(default=None) + max_clicks_settings: ProdegeMaxClicksSetting | None = Field(default=None) # This describes the project/surveygroup exclusions - past_participation: Optional[ProdegePastParticipation] = Field(default=None) + past_participation: ProdegePastParticipation | None = Field(default=None) # These describe the panelist exclusions/inclusions - include_psids: Optional[Set[UUIDStr]] = Field(default=None) - exclude_psids: Optional[Set[UUIDStr]] = Field(default=None) + include_psids: set[UUIDStr] | None = Field(default=None) + exclude_psids: set[UUIDStr] | None = Field(default=None) # There are no "qualifications" per se. Instead, everyone has the match a # parent quota (and its children) qualifications: List[str] = @@ -409,23 +409,23 @@ class ProdegeSurvey(MarketplaceTask): # The eligibility is somewhat complex, with parent and children quotas. # Going to keep it flat here. - quotas: List[ProdegeQuota] = Field(default_factory=list) + quotas: list[ProdegeQuota] = Field(default_factory=list) source: Literal[Source.PRODEGE] = Field(default=Source.PRODEGE) - used_question_ids: Set[ProdegeQuestionIdType] = Field(default_factory=set) + used_question_ids: set[ProdegeQuestionIdType] = 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, ProdegeCondition]] = Field(default=None) + conditions: dict[str, ProdegeCondition] = Field(default=None) # These do not come from the API. We set them. - created: Optional[AwareDatetimeISO] = Field( + created: AwareDatetimeISO | None = Field( description="when we created this survey in our system", default=None ) - updated: Optional[AwareDatetimeISO] = Field(default=None) + updated: AwareDatetimeISO | None = Field(default=None) @property def internal_id(self) -> str: @@ -494,7 +494,7 @@ class ProdegeSurvey(MarketplaceTask): return "1" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: ProdegeCondition( question_id="3", values=["1"], value_type=ConditionValueType.LIST @@ -510,7 +510,7 @@ class ProdegeSurvey(MarketplaceTask): return round(float(v), 2) @classmethod - def from_api(cls, d: Dict[str, Any]) -> Optional["ProdegeSurvey"]: + def from_api(cls, d: dict[str, Any]) -> "ProdegeSurvey" | None: try: return cls._from_api(d) except Exception as e: @@ -518,7 +518,7 @@ class ProdegeSurvey(MarketplaceTask): return None @classmethod - def _from_api(cls, d: Dict[str, Any]) -> "ProdegeSurvey": + def _from_api(cls, d: dict[str, Any]) -> "ProdegeSurvey": # Handle phases. keys in api response are 'loi' and 'actual_ir' if d["phases"]["loi_phase"] == "actual": @@ -575,14 +575,14 @@ class ProdegeSurvey(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)) return s @property - def quotas_verbose(self) -> List[List[Dict[str, Any]]]: + def quotas_verbose(self) -> list[list[dict[str, Any]]]: assert self.conditions is not None, "conditions must be set" res = [] for quota_group in self.quotas: @@ -625,7 +625,7 @@ class ProdegeSurvey(MarketplaceTask): return o1 == o2 - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump( mode="json", exclude={ @@ -655,7 +655,7 @@ class ProdegeSurvey(MarketplaceTask): return d @classmethod - def from_db(cls, d: Dict[str, Any]) -> "ProdegeSurvey": + 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["quotas"] = json.loads(d["quotas"]) @@ -675,7 +675,7 @@ class ProdegeSurvey(MarketplaceTask): def passes_quotas( self, - criteria_evaluation: Dict[str, Optional[bool]], + criteria_evaluation: dict[str, bool | None], country_iso: str, verbose: bool = False, ) -> bool: @@ -721,8 +721,8 @@ class ProdegeSurvey(MarketplaceTask): def passes_child_quotas( self, - criteria_evaluation: Dict[str, Optional[bool]], - child_quotas: List[ProdegeQuota], + criteria_evaluation: dict[str, bool | None], + child_quotas: list[ProdegeQuota], country_iso: str, verbose: bool = False, ) -> bool: @@ -753,14 +753,14 @@ class ProdegeSurvey(MarketplaceTask): return passes def determine_eligibility( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str + self, criteria_evaluation: dict[str, bool | None], country_iso: str ) -> bool: return self.is_open and self.passes_quotas( criteria_evaluation, country_iso=country_iso ) def print_eligibility( - self, criteria_evaluation: Dict[str, Optional[bool]], country_iso: str + self, criteria_evaluation: dict[str, bool | None], country_iso: str ) -> None: print(f"is_open: {self.is_open}") print("passes_quotas") @@ -769,5 +769,3 @@ class ProdegeSurvey(MarketplaceTask): criteria_evaluation, country_iso=country_iso, verbose=True ) ) - - return None diff --git a/generalresearch/models/prodege/task_collection.py b/generalresearch/models/prodege/task_collection.py index 0ea09e1..364c5ff 100644 --- a/generalresearch/models/prodege/task_collection.py +++ b/generalresearch/models/prodege/task_collection.py @@ -51,11 +51,11 @@ ProdegeTaskCollectionSchema = DataFrameSchema( class ProdegeTaskCollection(TaskCollection): - items: List[ProdegeSurvey] + items: list[ProdegeSurvey] _schema = ProdegeTaskCollectionSchema @staticmethod - def to_row(s: ProdegeSurvey) -> Dict[str, Any]: + def to_row(s: ProdegeSurvey) -> dict[str, Any]: fields = [ "survey_id", "status", diff --git a/generalresearch/models/repdata/question.py b/generalresearch/models/repdata/question.py index d33a4d8..bd5fded 100644 --- a/generalresearch/models/repdata/question.py +++ b/generalresearch/models/repdata/question.py @@ -4,7 +4,7 @@ import json import logging from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set +from typing import TYPE_CHECKING, Any, Literal from uuid import UUID from pydantic import ( @@ -35,12 +35,12 @@ class RepDataUserQuestionAnswer(BaseModel): # "anonymous" users, which are represented by a list of question answers # not associated with an actual user. No default b/c we must explicitly # set the field to None. - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32) + user_id: PositiveInt | None = Field(lt=MAX_INT32) question_id: str = Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") # This is optional b/c we do not need it when writing these to the db. When # these are fetched from the db for use in yield-management, we read this # field from the repdata_question table. - question_type: Optional[RepDataQuestionType] = Field(default=None) + question_type: RepDataQuestionType | None = Field(default=None) # 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]*$") @@ -55,10 +55,10 @@ class RepDataUserQuestionAnswer(BaseModel): ) @cached_property - def options_ids(self) -> Set[str]: + def options_ids(self) -> set[str]: return set(self.option_id.split("|")) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"question_type"}) d["created"] = self.created.replace(tzinfo=None) return d @@ -118,14 +118,14 @@ class RepDataQuestion(MarketplaceQuestion): question_name: str = Field( min_length=1, max_length=64, frozen=True, validation_alias="QualificationName" ) - lucid_id: Optional[str] = Field( + lucid_id: str | None = Field( min_length=1, max_length=16, pattern=r"^[0-9]+$", validation_alias="StandardGlobalID", frozen=True, ) - lucid_name: Optional[str] = Field( + lucid_name: str | None = Field( min_length=1, max_length=64, frozen=True, validation_alias="StandardGlobalName" ) question_text: str = Field( @@ -137,7 +137,7 @@ class RepDataQuestion(MarketplaceQuestion): question_type: RepDataQuestionType = Field( frozen=True, validation_alias="QualificationType" ) - options: Optional[List[RepDataQuestionOption]] = Field(default=None, min_length=1) + options: list[RepDataQuestionOption] | None = Field(default=None, min_length=1) source: Literal[Source.REPDATA] = Source.REPDATA @property @@ -160,8 +160,8 @@ class RepDataQuestion(MarketplaceQuestion): @classmethod def from_api( - cls, d: Dict[str, Any], country_iso: str, language_iso: str - ) -> Optional["RepDataQuestion"]: + cls, d: dict[str, Any], country_iso: str, language_iso: str + ) -> "RepDataQuestion" | None: """ :param d: Raw response from API """ @@ -173,7 +173,7 @@ class RepDataQuestion(MarketplaceQuestion): @classmethod def _from_api( - cls, d: Dict[str, Any], country_iso: str, language_iso: str + cls, d: dict[str, Any], country_iso: str, language_iso: str ) -> "RepDataQuestion": d["QualificationType"] = RepDataQuestionType.from_api(d["QualificationType"]) # zip code/age has a placeholder invalid option for some reason @@ -191,7 +191,7 @@ class RepDataQuestion(MarketplaceQuestion): ) @classmethod - def from_db(cls, d: Dict[str, Any]) -> "RepDataQuestion": + def from_db(cls, d: dict[str, Any]) -> "RepDataQuestion": options = None if d["options"]: options = [ @@ -212,7 +212,7 @@ class RepDataQuestion(MarketplaceQuestion): category_id=d.get("category_id"), ) - 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/repdata/survey.py b/generalresearch/models/repdata/survey.py index b6acd92..cfc58af 100644 --- a/generalresearch/models/repdata/survey.py +++ b/generalresearch/models/repdata/survey.py @@ -6,7 +6,7 @@ import logging from datetime import datetime, timezone from decimal import Decimal from functools import cached_property -from typing import Any, Dict, List, Literal, Optional, Set, Type +from typing import Any, Literal, Type from uuid import UUID from pydantic import ( @@ -54,13 +54,13 @@ class RepDataCondition(MarketplaceCondition): pattern=r"^[0-9]+$", validation_alias="StandardGlobalQuestionID", ) - values: List[str] = Field(min_length=1, validation_alias="PreCodes") + values: list[str] = Field(min_length=1, validation_alias="PreCodes") value_type: Literal[ConditionValueType.LIST] = Field( default=ConditionValueType.LIST ) @classmethod - def from_api(cls, d: Dict[str, Any]) -> "RepDataCondition": + def from_api(cls, d: dict[str, Any]) -> "RepDataCondition": if d["Condition"] == "Is": d["logical_operator"] = LogicalOperator.OR d["negate"] = False @@ -86,7 +86,7 @@ class RepDataQuota(BaseModel): ) quota_uuid: UUIDStr = Field(validation_alias="QuotaUd") name: str = Field(validation_alias="QuotaName") - desired_count: Optional[int] = Field( + desired_count: int | None = Field( default=None, validation_alias="Quota", description="Desired completes or starts (depending on calculation_type)", @@ -95,15 +95,15 @@ class RepDataQuota(BaseModel): validation_alias="QuotaAchieved", description="Achieved completes or starts (depending on calculation_type)", ) - remaining_count: Optional[int] = Field( + remaining_count: int | None = Field( validation_alias="QuotaRemaining", description="Completes or starts remaining (depending on calculation_type). Should " "be used as the indicator for whether more respondents are needed to a " "specific quota. If QuotaRemaining value = 0, then pause. If None, then the quota" "is completely open (i.e. infinity). Unclear if this is true though (see .is_open)", ) - conditions: List[RepDataCondition] = Field(min_length=1) - condition_hashes: List[str] = Field(min_length=1, default_factory=list) + conditions: list[RepDataCondition] = Field(min_length=1) + condition_hashes: list[str] = Field(min_length=1, default_factory=list) @field_validator("quota_uuid", mode="before") @classmethod @@ -137,7 +137,7 @@ class RepDataQuota(BaseModel): d = self.model_dump(mode="json", exclude={"conditions"}) return RepDataHashedQuota.model_validate(d) - def passes(self, criteria_evaluation: Dict[str, Optional[bool]]) -> Optional[bool]: + def passes(self, criteria_evaluation: dict[str, bool | None]) -> bool | None: # We have to match all conditions within the quota. return self.is_open and all( criteria_evaluation.get(c) for c in self.condition_hashes @@ -162,12 +162,12 @@ class RepDataStream(MarketplaceTask): validation_alias="CalculationType", ) - qualifications: List[RepDataCondition] = Field(min_length=1) - qualification_hashes: List[str] = Field(min_length=1, default_factory=list) - quotas: List[RepDataQuota] = Field(min_length=1) - hashed_quotas: List[RepDataHashedQuota] = Field(min_length=1, default_factory=list) + qualifications: list[RepDataCondition] = Field(min_length=1) + qualification_hashes: list[str] = Field(min_length=1, default_factory=list) + quotas: list[RepDataQuota] = Field(min_length=1) + hashed_quotas: list[RepDataHashedQuota] = Field(min_length=1, default_factory=list) - used_question_ids: Set[str] = Field(default_factory=set) + used_question_ids: set[str] = Field(default_factory=set) # Note: The API returns both Expected and ExpectedStreamCompletes which are the same expected_count: int = Field( @@ -182,7 +182,7 @@ class RepDataStream(MarketplaceTask): ) cpi: Decimal = Field(gt=0, le=100, validation_alias="CPI") - days_in_field: Optional[int] = Field(validation_alias="DaysInField", default=None) + days_in_field: int | None = Field(validation_alias="DaysInField", default=None) # # -------------- # # # Below here: these fields are useless because it is our own data. @@ -275,14 +275,14 @@ class RepDataStream(MarketplaceTask): @computed_field @cached_property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set(self.qualification_hashes.copy()) for q in self.hashed_quotas: s.update(set(q.condition_hashes)) return s @property - def all_conditions(self) -> List[RepDataCondition]: + def all_conditions(self) -> list[RepDataCondition]: cs = self.qualifications.copy() for quota in self.quotas: cs.extend(quota.conditions.copy()) @@ -312,7 +312,7 @@ class RepDataStream(MarketplaceTask): return "42" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: RepDataCondition( question_id="43", @@ -362,7 +362,7 @@ class RepDataStreamHashed(RepDataStream): @classmethod def from_db( - cls, res: Dict[str, Any], survey: RepDataSurveyHashed + cls, res: dict[str, Any], survey: RepDataSurveyHashed ) -> "RepDataStreamHashed": # We need certain fields copied over here so that a stream can exist # independent of the survey @@ -408,18 +408,18 @@ class RepDataSurvey(BaseModel): validation_alias="PII", description="Indicates whether PII is collected" ) - allowed_devices: List[DeviceType] = Field( + allowed_devices: list[DeviceType] = Field( min_length=1, validation_alias="Device Compatibility" ) - streams: List[RepDataStream] = Field(min_length=1) - hashed_streams: List[RepDataStreamHashed] = Field( + streams: list[RepDataStream] = Field(min_length=1) + hashed_streams: list[RepDataStreamHashed] = Field( min_length=1, default_factory=list ) # These do not come from the API. We set them ourselves - created: Optional[AwareDatetimeISO] = Field(default=None) - last_updated: Optional[AwareDatetimeISO] = Field(default=None) + created: AwareDatetimeISO | None = Field(default=None) + last_updated: AwareDatetimeISO | None = Field(default=None) @field_validator("survey_uuid", "project_uuid", mode="before") @classmethod @@ -434,7 +434,7 @@ class RepDataSurvey(BaseModel): return data @field_validator("allowed_devices", mode="after") - def sort_allowed_devices(cls, values: List[str]): + def sort_allowed_devices(cls, values: list[str]): return sorted(values) @property @@ -452,14 +452,14 @@ class RepDataSurvey(BaseModel): return self.survey_status == RepDataStatus.LIVE @property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set() for stream in self.hashed_streams: s.update(stream.all_hashes) return s @property - def all_conditions(self) -> List[RepDataCondition]: + def all_conditions(self) -> list[RepDataCondition]: cs = list() for stream in self.streams: cs.extend(stream.all_conditions) @@ -472,7 +472,7 @@ class RepDataSurvey(BaseModel): return ",".join(map(str, sorted([d.value for d in self.allowed_devices]))) @classmethod - def from_api(cls, survey_response) -> Optional["RepDataSurvey"]: + def from_api(cls, survey_response) -> "RepDataSurvey" | None: """ :param survey_response: Raw response from API """ @@ -518,7 +518,7 @@ class RepDataSurvey(BaseModel): def is_changed(self, other) -> bool: return not self.is_unchanged(other) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: return self.to_hashed_survey().to_mysql() def to_hashed_survey(self) -> "RepDataSurveyHashed": @@ -533,7 +533,7 @@ class RepDataSurveyHashed(RepDataSurvey): streams: None = Field(default=None, exclude=True) @classmethod - def from_db(cls, res: Dict[str, Any]) -> "RepDataSurveyHashed": + def from_db(cls, res: dict[str, Any]) -> "RepDataSurveyHashed": res["allowed_devices"] = [ DeviceType(int(x)) for x in res["allowed_devices"].split(",") ] @@ -542,7 +542,7 @@ class RepDataSurveyHashed(RepDataSurvey): res["last_updated"] = res["last_updated"].replace(tzinfo=timezone.utc) return cls.model_validate(res) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", by_alias=True, exclude={"hashed_streams"}) d["allowed_devices"] = ",".join( map(str, sorted([d.value for d in self.allowed_devices])) diff --git a/generalresearch/models/repdata/task_collection.py b/generalresearch/models/repdata/task_collection.py index 740670b..b79808d 100644 --- a/generalresearch/models/repdata/task_collection.py +++ b/generalresearch/models/repdata/task_collection.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, List +from __future__ import annotations + +from typing import Any import pandas as pd from pandera import Check, Column, DataFrameSchema, Index @@ -46,8 +48,8 @@ RepDataTaskCollectionSchema = DataFrameSchema( "remaining_count": Column(int, Check.greater_than_or_equal_to(0)), "calculation_type": Column(str, Check.isin(TaskCalculationType)), "cpi": Column(float, Check.between(min_value=0, max_value=100)), - "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=[ # # There's only 1 or 2 live surveys, so we can't really assert anything ... @@ -78,10 +80,10 @@ RepDataTaskCollectionSchema = DataFrameSchema( class RepDataTaskCollection(TaskCollection): - items: List[RepDataSurveyHashed] + items: list[RepDataSurveyHashed] _schema = RepDataTaskCollectionSchema - def to_rows(self, s: RepDataSurveyHashed) -> List[Dict[str, Any]]: + def to_rows(self, s: RepDataSurveyHashed) -> list[dict[str, Any]]: survey_fields = [ "survey_id", "survey_uuid", diff --git a/generalresearch/models/sago/question.py b/generalresearch/models/sago/question.py index 62ff363..f1ef8b4 100644 --- a/generalresearch/models/sago/question.py +++ b/generalresearch/models/sago/question.py @@ -1,10 +1,12 @@ +from __future__ import annotations + # https://developer-beta.market-cube.com/api-details#api=definition-api&operation=get-api-v1-definition-qualification # -answers-lanaguge-languageid import json import logging from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set +from typing import TYPE_CHECKING, Any, Literal from pydantic import ( BaseModel, @@ -39,7 +41,7 @@ class SagoQuestionOption(BaseModel): ) # This is returned by the API but does not seem to be used for anything. # Will keep it any ways. - code: Optional[str] = Field(min_length=1, max_length=16) + code: str | None = Field(min_length=1, max_length=16) text: str = Field( min_length=1, max_length=1024, @@ -94,13 +96,13 @@ class SagoUserQuestionAnswer(BaseModel): # "anonymous" users, which are represented by a list of question answers # not associated with an actual user. No default b/c we must explicitly set # the field to None. - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32) + user_id: PositiveInt | None = Field(lt=MAX_INT32) question_id: str = Field(min_length=1, max_length=16, pattern=r"^[0-9]+$") # This is optional b/c we do not need it when writing these to the db. When # these are fetched from the db for use in yield-management, we read this # field from the question table. - question_type: Optional[SagoQuestionType] = Field(default=None) + question_type: SagoQuestionType | None = Field(default=None) # This may be a pipe-separated string if the question_type is multi. # regex means any chars except capital letters @@ -121,10 +123,10 @@ class SagoUserQuestionAnswer(BaseModel): return self.pre_code @cached_property - def options_ids(self) -> Set[str]: + def options_ids(self) -> set[str]: return set(self.pre_code.split("|")) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"question_type"}) d["created"] = self.created.replace(tzinfo=None) return d @@ -145,10 +147,10 @@ class SagoQuestion(MarketplaceQuestion): max_length=1024, min_length=1, description="The text shown to respondents" ) question_type: SagoQuestionType = Field(frozen=True) - options: Optional[List[SagoQuestionOption]] = Field(default=None, min_length=1) + options: list[SagoQuestionOption] | None = Field(default=None, min_length=1) # This comes from the API field "qualificationCategoryId" - tags: Optional[str] = Field(default=None, frozen=True) + tags: str | None = Field(default=None, frozen=True) source: Literal[Source.SAGO] = Source.SAGO @property @@ -165,13 +167,13 @@ class SagoQuestion(MarketplaceQuestion): return self @field_validator("question_name", "question_text", "tags", mode="after") - def remove_nbsp(cls, s: Optional[str]): + def remove_nbsp(cls, s: str | None): return string_utils.remove_nbsp(s) @classmethod def from_api( - cls, d: Dict[str, Any], country_iso: str, language_iso: str - ) -> Optional["SagoQuestion"]: + cls, d: dict[str, Any], country_iso: str, language_iso: str + ) -> "SagoQuestion" | None: """ :param d: Raw response from API :param country_iso: @@ -186,7 +188,7 @@ class SagoQuestion(MarketplaceQuestion): @classmethod def _from_api( - cls, d: Dict[str, Any], country_iso: str, language_iso: str + cls, d: dict[str, Any], country_iso: str, language_iso: str ) -> "SagoQuestion": sago_category_to_tags = { 1: "Standard", @@ -221,7 +223,7 @@ class SagoQuestion(MarketplaceQuestion): ) @classmethod - def from_db(cls, d: Dict[str, Any]) -> "SagoQuestion": + def from_db(cls, d: dict[str, Any]) -> "SagoQuestion": options = None if d["options"]: options = [ @@ -243,7 +245,7 @@ class SagoQuestion(MarketplaceQuestion): category_id=d.get("category_id"), ) - 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/sago/survey.py b/generalresearch/models/sago/survey.py index a8188b0..b07c846 100644 --- a/generalresearch/models/sago/survey.py +++ b/generalresearch/models/sago/survey.py @@ -5,7 +5,7 @@ import logging from datetime import timezone from decimal import Decimal from functools import cached_property -from typing import Annotated, Any, Dict, List, Literal, Optional, Set, Tuple, Type +from typing import Annotated, Any, Literal, Type from more_itertools import flatten from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator @@ -38,18 +38,18 @@ locale_helper = Localelator() class SagoCondition(MarketplaceCondition): model_config = ConfigDict(populate_by_name=True, frozen=False, extra="ignore") - question_id: Optional[CoercedStr] = Field( + question_id: CoercedStr | None = Field( min_length=1, max_length=16, pattern=r"^[0-9]+$" ) # There isn't really a hard limit, but their API is inconsistent and # sometimes returns all the options comma-separated instead of as a list. # Try to catch that. - values: List[Annotated[str, Field(max_length=128)]] = Field() + values: list[Annotated[str, Field(max_length=128)]] = Field() _CONVERT_LIST_TO_RANGE = ["59"] @classmethod - def from_api(cls, d: Dict[str, Any]) -> "SagoCondition": + def from_api(cls, d: dict[str, Any]) -> "SagoCondition": d["logical_operator"] = LogicalOperator.OR d["value_type"] = ConditionValueType(d["value_type"]) d["negate"] = False @@ -67,7 +67,7 @@ class SagoQuota(BaseModel): cpi: Decimal = Field(gt=0, le=100, decimal_places=2, max_digits=5) remaining_count: int = Field() - condition_hashes: List[str] = Field(min_length=0, default_factory=list) + condition_hashes: list[str] = Field(min_length=0, default_factory=list) # There is no explicit status. The quota is closed if the count is 0 @@ -80,23 +80,23 @@ class SagoQuota(BaseModel): return self.remaining_count >= min_open_spots @classmethod - def from_api(cls, d: Dict[str, Any]) -> Self: + def from_api(cls, d: dict[str, Any]) -> Self: return cls.model_validate(d) - 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) - 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. # We can "match" a quota that is closed. In that case, we would not be # eligible for the survey. return all(criteria_evaluation.get(c) for c in self.condition_hashes) def matches_optional( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Optional[bool]: + self, criteria_evaluation: dict[str, bool | None] + ) -> bool | None: # We need to know if any conditions are unknown to avoid matching a # full quota. If any fail, then we know we fail regardless of any # being unknown. @@ -108,8 +108,8 @@ class SagoQuota(BaseModel): return True def matches_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]]: # Passes back "matches" (T/F/none) and a list of unknown criterion hashes hash_evals = { cell: criteria_evaluation.get(cell) for cell in self.condition_hashes @@ -151,35 +151,35 @@ class SagoSurvey(MarketplaceTask): allowed_devices: DeviceTypes = Field(min_length=1) collects_pii: bool = Field(default=False) - survey_exclusions: Optional[AlphaNumStrSet] = Field( + survey_exclusions: AlphaNumStrSet | None = Field( description="list of excluded survey ids", default=None ) - ip_exclusions: Optional[IPLikeStrSet] = Field( + ip_exclusions: IPLikeStrSet | None = Field( description="list of excluded IP addresses", default=None ) # Documentation I think is wrong. These are the keys "LOI" and "IR". it # doesn't say that they are bid or not, but they never seem to change ... - bid_loi: Optional[int] = Field(default=None, le=120 * 60) - bid_ir: Optional[float] = Field(default=None, ge=0, le=1) + bid_loi: int | None = Field(default=None, le=120 * 60) + bid_ir: float | None = Field(default=None, ge=0, le=1) live_link: str = Field() # this comes from the Survey Reservation endpoint remaining_count: int = Field() - qualifications: List[str] = Field(default_factory=list) - quotas: List[SagoQuota] = Field(default_factory=list) + qualifications: list[str] = Field(default_factory=list) + quotas: list[SagoQuota] = Field(default_factory=list) source: Literal[Source.SAGO] = Field(default=Source.SAGO) - used_question_ids: Set[AlphaNumStr] = Field(default_factory=set) + used_question_ids: set[AlphaNumStr] = 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, SagoCondition]] = Field(default=None) + conditions: dict[str, SagoCondition] | None = Field(default=None) # These come from the API modified_api: AwareDatetimeISO = Field( @@ -187,8 +187,8 @@ class SagoSurvey(MarketplaceTask): ) # 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: @@ -208,7 +208,7 @@ class SagoSurvey(MarketplaceTask): @computed_field @cached_property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set(self.qualifications) for q in self.quotas: s.update(set(q.condition_hashes)) @@ -243,7 +243,7 @@ class SagoSurvey(MarketplaceTask): return "59" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: SagoCondition( question_id="60", @@ -259,7 +259,7 @@ class SagoSurvey(MarketplaceTask): } @classmethod - def from_api(cls, d: Dict[str, Any]) -> Optional["SagoSurvey"]: + def from_api(cls, d: dict[str, Any]) -> "SagoSurvey" | None: try: return cls._from_api(d) except Exception as e: @@ -267,7 +267,7 @@ class SagoSurvey(MarketplaceTask): return None @classmethod - def _from_api(cls, d: Dict[str, Any]) -> "SagoSurvey": + def _from_api(cls, d: dict[str, Any]) -> "SagoSurvey": return cls.model_validate(d) def __repr__(self) -> str: @@ -294,7 +294,7 @@ class SagoSurvey(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={ @@ -313,7 +313,7 @@ class SagoSurvey(MarketplaceTask): return d @classmethod - def from_db(cls, d: Dict[str, Any]): + 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) @@ -323,14 +323,14 @@ class SagoSurvey(MarketplaceTask): return cls.model_validate(d) def passes_qualifications( - self, criteria_evaluation: Dict[str, Optional[bool]] + self, criteria_evaluation: dict[str, bool | None] ) -> bool: # We have to match all quals return all(criteria_evaluation.get(q) for q in self.qualifications) def passes_qualifications_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]]: # Passes back "passes" (T/F/none) and a list of unknown criterion hashes hash_evals = {q: criteria_evaluation.get(q) for q in self.qualifications} evals = set(hash_evals.values()) @@ -342,7 +342,7 @@ class SagoSurvey(MarketplaceTask): return None, {cell for cell, ev in hash_evals.items() if ev is None} return True, set() - def passes_quotas(self, criteria_evaluation: Dict[str, Optional[bool]]) -> bool: + def passes_quotas(self, criteria_evaluation: dict[str, bool | None]) -> bool: # Many surveys have 0 quotas. Quotas are exclusionary. # They can NOT match a quota where currently_open=0 any_pass = True @@ -354,8 +354,8 @@ class SagoSurvey(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]]: # Many surveys have 0 quotas. Quotas are exclusionary. # They can NOT match a quota where currently_open=0 if len(self.quotas) == 0: @@ -394,7 +394,7 @@ class SagoSurvey(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 @@ -403,8 +403,8 @@ class SagoSurvey(MarketplaceTask): ) def determine_eligibility_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]]: if self.is_open is False: return False, set() pass_quals, h_quals = self.passes_qualifications_soft(criteria_evaluation) diff --git a/generalresearch/models/sago/task_collection.py b/generalresearch/models/sago/task_collection.py index 41d047e..08de55f 100644 --- a/generalresearch/models/sago/task_collection.py +++ b/generalresearch/models/sago/task_collection.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, List, Set +from __future__ import annotations + +from typing import Any import pandas as pd from pandera import Check, Column, DataFrameSchema, Index @@ -11,8 +13,8 @@ from generalresearch.models.thl.survey.task_collection import ( create_empty_df_from_schema, ) -COUNTRY_ISOS: Set[str] = Localelator().get_all_countries() -LANGUAGE_ISOS: Set[str] = Localelator().get_all_languages() +COUNTRY_ISOS: set[str] = Localelator().get_all_countries() +LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() SagoTaskCollectionSchema = DataFrameSchema( columns={ @@ -31,8 +33,8 @@ SagoTaskCollectionSchema = DataFrameSchema( "remaining_count": Column(int), "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( @@ -48,10 +50,10 @@ SagoTaskCollectionSchema = DataFrameSchema( class SagoTaskCollection(TaskCollection): - items: List[SagoSurvey] + items: list[SagoSurvey] _schema = SagoTaskCollectionSchema - def to_row(self, s: SagoSurvey) -> Dict[str, Any]: + def to_row(self, s: SagoSurvey) -> dict[str, Any]: d = s.model_dump( mode="json", exclude={ diff --git a/generalresearch/models/spectrum/question.py b/generalresearch/models/spectrum/question.py index 4f8b5e1..b544c5e 100644 --- a/generalresearch/models/spectrum/question.py +++ b/generalresearch/models/spectrum/question.py @@ -6,7 +6,7 @@ import logging from datetime import datetime, timezone from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set +from typing import TYPE_CHECKING, Any, Literal from uuid import UUID from pydantic import ( @@ -41,12 +41,12 @@ class SpectrumUserQuestionAnswer(BaseModel): # answers not associated with an actual user. No default b/c we must # explicitly set the field to None. - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32) + user_id: PositiveInt | None = Field(lt=MAX_INT32) question_id: SpectrumQuestionIdType = Field() # This is optional b/c we do not need it when writing these to the # db. When these are fetched from the db for use in yield-management, # we read this field from the spectrum_question table. - question_type: Optional[SpectrumQuestionType] = Field(default=None) + question_type: SpectrumQuestionType | None = Field(default=None) # 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]*$") @@ -63,10 +63,10 @@ class SpectrumUserQuestionAnswer(BaseModel): ) @cached_property - def options_ids(self) -> Set[str]: + def options_ids(self) -> set[str]: return set(self.option_id.split("|")) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"question_type"}) d["created"] = self.created.replace(tzinfo=None) return d @@ -123,7 +123,7 @@ class SpectrumQuestionType(str, Enum): UNKNOWN = "u" @staticmethod - def get_api_map() -> Dict[int, SpectrumQuestionType]: + def get_api_map() -> dict[int, SpectrumQuestionType]: return { 1: SpectrumQuestionType.SINGLE_SELECT, 2: SpectrumQuestionType.SINGLE_SELECT, @@ -169,8 +169,8 @@ class SpectrumQuestion(MarketplaceQuestion): # looks to be a comma-separated str of "tags" or keywords associated # with a question, but they are freeform and don't pertain to any sort # of structured schema. This will be useful ChatGPT - tags: Optional[str] = Field(default=None, frozen=True) - options: Optional[List[SpectrumQuestionOption]] = Field( + tags: str | None = Field(default=None, frozen=True) + options: list[SpectrumQuestionOption] | None = Field( default=None, min_length=1, frozen=True ) # This comes from the API. Of course there are more than what is documented. @@ -178,7 +178,7 @@ class SpectrumQuestion(MarketplaceQuestion): class_num: SpectrumQuestionClass = Field(frozen=True) # This comes from the API. It is when it was created in Spectrum's DB, # not when we created it - created: Optional[AwareDatetimeISO] = Field(default=None, frozen=True) + created: AwareDatetimeISO | None = Field(default=None, frozen=True) source: Literal[Source.SPECTRUM] = Source.SPECTRUM @@ -207,7 +207,7 @@ class SpectrumQuestion(MarketplaceQuestion): return data @field_validator("question_name", "question_text", "tags", mode="after") - def remove_nbsp(cls, s: Optional[str]): + def remove_nbsp(cls, s: str | None): return string_utils.remove_nbsp(s) @model_validator(mode="before") @@ -236,7 +236,7 @@ class SpectrumQuestion(MarketplaceQuestion): @field_validator("options") @classmethod - def uniquify_options(cls, options: Optional[List[SpectrumQuestionOption]]): + def uniquify_options(cls, options: list[SpectrumQuestionOption] | None): if options: # The API returns questions with identical option IDs multiple # times. They seem to all be typo/corrections to the text, so @@ -251,8 +251,8 @@ class SpectrumQuestion(MarketplaceQuestion): @classmethod def from_api( - cls, d: Dict[str, Any], country_iso: str, language_iso: str - ) -> Optional["SpectrumQuestion"]: + cls, d: dict[str, Any], country_iso: str, language_iso: str + ) -> "SpectrumQuestion" | None: # To not pollute our logs, we know we are skipping any question that # meets the following conditions: if not SpectrumQuestionType.from_api(d["type"]): @@ -268,7 +268,7 @@ class SpectrumQuestion(MarketplaceQuestion): return None @classmethod - def _from_api(cls, d: Dict[str, Any], country_iso: str, language_iso: str) -> Self: + def _from_api(cls, d: dict[str, Any], country_iso: str, language_iso: str) -> Self: options = None if d.get("condition_codes"): # Sometimes they use the key "name" instead of "text" ... ? @@ -301,7 +301,7 @@ class SpectrumQuestion(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 = [ @@ -329,7 +329,7 @@ class SpectrumQuestion(MarketplaceQuestion): created=d["created"], ) - 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"]) if self.created: diff --git a/generalresearch/models/spectrum/survey.py b/generalresearch/models/spectrum/survey.py index 72093eb..3683745 100644 --- a/generalresearch/models/spectrum/survey.py +++ b/generalresearch/models/spectrum/survey.py @@ -4,7 +4,7 @@ import json import logging from datetime import timezone from decimal import Decimal -from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Type +from typing import Any, Literal, Type from more_itertools import flatten from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator @@ -37,7 +37,7 @@ locale_helper = Localelator() class SpectrumCondition(MarketplaceCondition): model_config = ConfigDict(populate_by_name=True, frozen=False, extra="ignore") - question_id: Optional[CoercedStr] = Field( + question_id: CoercedStr | None = Field( min_length=1, max_length=16, pattern=r"^[0-9]+$", @@ -65,7 +65,7 @@ class SpectrumCondition(MarketplaceCondition): return self @classmethod - def from_api(cls, d: Dict[str, Any]) -> "SpectrumCondition": + def from_api(cls, d: dict[str, Any]) -> "SpectrumCondition": """Ranges can get returns with a key "units" indicating years or months. This is ridiculous, and we don't ask for birthdate, so we can't really get month accuracy. Normalize to years. @@ -98,7 +98,7 @@ class SpectrumQuota(BaseModel): description="Number of completes currently available in the quota. If " "the value is 0, any respondent matching this quota will be rejected." ) - condition_hashes: List[str] = Field(min_length=0, default_factory=list) + condition_hashes: list[str] = Field(min_length=0, default_factory=list) # API also returns remaining & achieved, but these are supplier-scoped. # There is no explicit status. The quota is closed if the count is 0 @@ -114,23 +114,23 @@ class SpectrumQuota(BaseModel): return self.remaining_count >= min_open_spots @classmethod - def from_api(cls, d: Dict) -> Self: + def from_api(cls, d: dict) -> Self: d["remaining_count"] = d["quantities"]["currently_open"] return cls.model_validate(d) - 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) - 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. We can "match" a quota that is # closed. In that case, we would not be eligible for the survey. return all(criteria_evaluation.get(c) for c in self.condition_hashes) def matches_optional( - self, criteria_evaluation: Dict[str, Optional[bool]] - ) -> Optional[bool]: + self, criteria_evaluation: dict[str, bool | None] + ) -> bool | None: # We need to know if any conditions are unknown to avoid matching a # full quota. If any fail, then we know we fail regardless of any # being unknown. @@ -142,8 +142,8 @@ class SpectrumQuota(BaseModel): return True def matches_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]]: # Passes back "matches" (T/F/none) and a list of unknown criterion hashes hash_evals = { cell: criteria_evaluation.get(cell) for cell in self.condition_hashes @@ -188,7 +188,7 @@ class SpectrumSurvey(MarketplaceTask): # called "survey_grouping" in API. If a respondent has previously taken any # of these surveys, they will be excluded if that survey was taken in # the exclusion_period. - survey_exclusions: Optional[AlphaNumStrSet] = Field( + survey_exclusions: AlphaNumStrSet | None = Field( description="list of excluded survey ids", default=None ) exclusion_period: int = Field(default=30, description="in days") @@ -200,16 +200,16 @@ class SpectrumSurvey(MarketplaceTask): # when we update the db, we must not overwrite this with NULL. # API key: "survey_performance" - bid_loi: Optional[int] = Field(default=None, le=120 * 60) - bid_ir: Optional[float] = Field(default=None, ge=0, le=1) - overall_loi: Optional[int] = Field(default=None, le=120 * 60) - overall_ir: Optional[float] = Field(default=None, ge=0, le=1) - last_block_loi: Optional[int] = Field(default=None, le=120 * 60) - last_block_ir: Optional[float] = Field(default=None, ge=0, le=1) + bid_loi: int | None = Field(default=None, le=120 * 60) + bid_ir: float | None = Field(default=None, ge=0, le=1) + overall_loi: int | None = Field(default=None, le=120 * 60) + overall_ir: float | None = Field(default=None, ge=0, le=1) + last_block_loi: int | None = Field(default=None, le=120 * 60) + last_block_ir: float | None = Field(default=None, ge=0, le=1) # Undocumented. They sent us an email indicating that this is the last time # there was a complete for all suppliers on this survey. - project_last_complete_date: Optional[AwareDatetimeISO] = Field(default=None) + project_last_complete_date: AwareDatetimeISO | None = Field(default=None) # ISO 3166-1 alpha-2 (two-letter codes, lowercase) country_iso: str = Field( @@ -224,21 +224,21 @@ class SpectrumSurvey(MarketplaceTask): # list is an exclude or include list. If incl_excl = 1, the survey has an # include list, and only those in the list are eligible. This list gets # updated everytime someone on the list takes the survey. - include_psids: Optional[UUIDStrSet] = Field(default=None) - exclude_psids: Optional[UUIDStrSet] = Field(default=None) + include_psids: UUIDStrSet | None = Field(default=None) + exclude_psids: UUIDStrSet | None = Field(default=None) - qualifications: List[str] = Field(default_factory=list) - quotas: List[SpectrumQuota] = Field(default_factory=list) + qualifications: list[str] = Field(default_factory=list) + quotas: list[SpectrumQuota] = Field(default_factory=list) source: Literal[Source.SPECTRUM] = Field(default=Source.SPECTRUM) - used_question_ids: Set[AlphaNumStr] = Field(default_factory=set) + used_question_ids: set[AlphaNumStr] = 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, SpectrumCondition]] = Field(default=None) + conditions: dict[str, SpectrumCondition] | None = Field(default=None) # These come from the API created_api: AwareDatetimeISO = Field( @@ -250,7 +250,7 @@ class SpectrumSurvey(MarketplaceTask): ) # This does not come from the API. We set it when we update this in the db. - updated: Optional[AwareDatetimeISO] = Field(default=None) + updated: AwareDatetimeISO | None = Field(default=None) @property def internal_id(self) -> str: @@ -270,7 +270,7 @@ class SpectrumSurvey(MarketplaceTask): @computed_field @property - def all_hashes(self) -> Set[str]: + def all_hashes(self) -> set[str]: s = set(self.qualifications) for q in self.quotas: s.update(set(q.condition_hashes)) @@ -305,7 +305,7 @@ class SpectrumSurvey(MarketplaceTask): return "212" @property - def marketplace_genders(self) -> Dict[Gender, Optional[MarketplaceCondition]]: + def marketplace_genders(self) -> dict[Gender, MarketplaceCondition | None]: return { Gender.MALE: SpectrumCondition( question_id="211", @@ -321,7 +321,7 @@ class SpectrumSurvey(MarketplaceTask): } @classmethod - def from_api(cls, d: Dict[str, Any]) -> Optional["SpectrumSurvey"]: + def from_api(cls, d: dict[str, Any]) -> "SpectrumSurvey" | None: try: return cls._from_api(d) except Exception as e: @@ -329,7 +329,7 @@ class SpectrumSurvey(MarketplaceTask): return None @classmethod - def _from_api(cls, d: Dict[str, Any]) -> Self: + def _from_api(cls, d: dict[str, Any]) -> Self: assert d["click_balancing"] in {0, 1}, "unknown click_balancing value" d["calculation_type"] = ( TaskCalculationType.STARTS @@ -366,7 +366,7 @@ class SpectrumSurvey(MarketplaceTask): exclude={"updated", "conditions"} ) - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump( mode="json", exclude={ @@ -388,7 +388,7 @@ class SpectrumSurvey(MarketplaceTask): return d @classmethod - def from_db(cls, d: Dict[str, Any]) -> Self: + 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) @@ -415,14 +415,14 @@ class SpectrumSurvey(MarketplaceTask): """ def passes_qualifications( - self, criteria_evaluation: Dict[str, Optional[bool]] + self, criteria_evaluation: dict[str, bool | None] ) -> bool: # We have to match all quals return all(criteria_evaluation.get(q) for q in self.qualifications) def passes_qualifications_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]]: # Passes back "passes" (T/F/none) and a list of unknown criterion hashes hash_evals = {q: criteria_evaluation.get(q) for q in self.qualifications} # We have to match all. So if any are False, we know we don't pass @@ -435,7 +435,7 @@ class SpectrumSurvey(MarketplaceTask): return None, {cell for cell, ev in hash_evals.items() if ev is None} return True, set() - 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 at least 1 quota, but they can NOT match a quota # where currently_open=0 any_pass = False @@ -450,8 +450,8 @@ class SpectrumSurvey(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]]: # We have to match at least 1 quota, but they can NOT match a quota # where currently_open=0 quota_eval = { @@ -488,7 +488,7 @@ class SpectrumSurvey(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 @@ -497,8 +497,8 @@ class SpectrumSurvey(MarketplaceTask): ) def determine_eligibility_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]]: if self.is_open is False: return False, set() pass_quals, h_quals = self.passes_qualifications_soft(criteria_evaluation) diff --git a/generalresearch/models/spectrum/task_collection.py b/generalresearch/models/spectrum/task_collection.py index 3c2ee85..4db866b 100644 --- a/generalresearch/models/spectrum/task_collection.py +++ b/generalresearch/models/spectrum/task_collection.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Set +from __future__ import annotations import pandas as pd from pandera import Check, Column, DataFrameSchema, Index @@ -12,8 +12,8 @@ from generalresearch.models.thl.survey.task_collection import ( create_empty_df_from_schema, ) -COUNTRY_ISOS: Set[str] = Localelator().get_all_countries() -LANGUAGE_ISOS: Set[str] = Localelator().get_all_languages() +COUNTRY_ISOS: set[str] = Localelator().get_all_countries() +LANGUAGE_ISOS: set[str] = Localelator().get_all_languages() SpectrumTaskCollectionSchema = DataFrameSchema( columns={ @@ -43,8 +43,8 @@ SpectrumTaskCollectionSchema = DataFrameSchema( "created_api": Column(dtype=pd.DatetimeTZDtype(tz="UTC")), "modified_api": 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( @@ -60,10 +60,10 @@ SpectrumTaskCollectionSchema = DataFrameSchema( class SpectrumTaskCollection(TaskCollection): - items: List[SpectrumSurvey] + items: list[SpectrumSurvey] _schema = SpectrumTaskCollectionSchema - def to_rows(self, s: SpectrumSurvey) -> List[Dict]: + def to_rows(self, s: SpectrumSurvey) -> list[dict]: fields = [ "survey_name", "status", diff --git a/generalresearch/models/thl/__init__.py b/generalresearch/models/thl/__init__.py index 0356842..abc129b 100644 --- a/generalresearch/models/thl/__init__.py +++ b/generalresearch/models/thl/__init__.py @@ -1,5 +1,4 @@ from decimal import Decimal -from typing import Optional from generalresearch.models.thl.finance import ( POPFinancial, @@ -24,11 +23,11 @@ PayoutEvent.model_rebuild() BrokerageProductPayoutEvent.model_rebuild() -def decimal_to_int_cents(usd: Optional[Decimal]) -> Optional[int]: +def decimal_to_int_cents(usd: Decimal | None) -> int | None: return round(usd * 100) if usd is not None else None -def int_cents_to_decimal(value: Optional[int], decimals: int = 2) -> Optional[Decimal]: +def int_cents_to_decimal(value: int | None, decimals: int = 2) -> Decimal | None: if value is None: return None return (Decimal(value) / Decimal(100)).quantize(Decimal(10) ** -decimals) diff --git a/generalresearch/models/thl/category.py b/generalresearch/models/thl/category.py index 4e9e2ff..1ed436a 100644 --- a/generalresearch/models/thl/category.py +++ b/generalresearch/models/thl/category.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, Optional +from __future__ import annotations + +from typing import Any from uuid import uuid4 from pydantic import BaseModel, Field, PositiveInt, model_validator @@ -8,11 +10,11 @@ from generalresearch.models.custom_types import UUIDStr class Category(BaseModel, frozen=True): - id: Optional[PositiveInt] = Field(exclude=True, default=None) + id: PositiveInt | None = Field(exclude=True, default=None) uuid: UUIDStr = Field(examples=[uuid4().hex]) - adwords_vertical_id: Optional[str] = Field(default=None, max_length=8) + adwords_vertical_id: str | None = Field(default=None, max_length=8) label: str = Field(max_length=255, examples=["Hair Loss"]) @@ -23,8 +25,8 @@ class Category(BaseModel, frozen=True): examples=["/Beauty & Fitness/Hair Care/Hair Loss"], ) - parent_id: Optional[PositiveInt] = Field(default=None, exclude=True) - parent_uuid: Optional[UUIDStr] = Field(default=None, examples=[uuid4().hex]) + parent_id: PositiveInt | None = Field(default=None, exclude=True) + parent_uuid: UUIDStr | None = Field(default=None, examples=[uuid4().hex]) @model_validator(mode="after") def check_path(self) -> Self: @@ -45,7 +47,7 @@ class Category(BaseModel, frozen=True): return self.path.split("/", 2)[1] @property - def parent_path(self) -> Optional[str]: + def parent_path(self) -> str | None: # If path is "/Beauty & Fitness/Hair Care/Hair Loss", this returns "/Beauty & Fitness/Hair Care" return self.path.rsplit("/", 1)[0] or None @@ -53,7 +55,7 @@ class Category(BaseModel, frozen=True): def is_root(self) -> bool: return self.parent_path is None - def to_offerwall_api(self) -> Dict[str, Any]: + def to_offerwall_api(self) -> dict[str, Any]: return { "id": self.uuid, "label": self.label, diff --git a/generalresearch/models/thl/contest/__init__.py b/generalresearch/models/thl/contest/__init__.py index f9693eb..363c8c0 100644 --- a/generalresearch/models/thl/contest/__init__.py +++ b/generalresearch/models/thl/contest/__init__.py @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import datetime, timezone -from typing import Any, Dict, Optional from uuid import uuid4 from pydantic import ( @@ -24,12 +23,12 @@ class ContestEntryRule(BaseModel): Only applies if the ContestType is ENTRY! """ - max_entry_amount_per_user: Optional[USDCent | PositiveInt] = Field( + max_entry_amount_per_user: USDCent | PositiveInt | None = Field( description="Maximum total value of entries per user", default=None, ) - max_daily_entries_per_user: Optional[PositiveInt] = Field( + max_daily_entries_per_user: PositiveInt | None = Field( description="Maximum entries per user allowed per day for this contest", default=None, ) @@ -37,9 +36,9 @@ class ContestEntryRule(BaseModel): # TODO: Only allow entries if user meets some criteria: gold-membership # status, ID/phone verified, min_completes etc... Maybe these get put # in a separate model b/c the could apply if the ContestType is not ENTRY - min_completes: Optional[int] = None - min_membership_level: Optional[int] = None - id_verified: Optional[bool] = None + min_completes: int | None = None + min_membership_level: int | None = None + id_verified: bool | None = None class ContestEndCondition(BaseModel): @@ -54,7 +53,7 @@ class ContestEndCondition(BaseModel): description="The contest is over once this amount is reached. (sum of all entry amount)", ) # In a LeaderboardContest, ends_at equals the leaderboard's end period plus 90 minutes - ends_at: Optional[AwareDatetimeISO] = Field( + ends_at: AwareDatetimeISO | None = Field( default=None, description="The contest is over at this time." ) @@ -63,22 +62,22 @@ class ContestPrize(BaseModel): kind: ContestPrizeKind = Field( description=ContestPrizeKind.as_openapi_with_value_descriptions() ) - name: Optional[str] = Field(default=None) - description: Optional[str] = Field(default=None) + name: str | None = Field(default=None) + description: str | None = Field(default=None) estimated_cash_value: USDCent = Field( description="Estimated cash value of prize in USDCents", ) - cash_amount: Optional[USDCent] = Field( + cash_amount: USDCent | None = Field( default=None, description="If the kind=ContestPrizeKind.CASH, this is the amount of the prize", ) - promotion_id: Optional[UUIDStr] = Field( + promotion_id: UUIDStr | None = Field( default=None, description="If the kind=ContestPrizeKind.PROMOTION, this is the promotion ID", ) # only if the contest.contest_type = LEADERBOARD - leaderboard_rank: Optional[PositiveInt] = Field( + leaderboard_rank: PositiveInt | None = Field( default=None, description="The prize is for achieving this rank in the associated " "leaderboard. The highest rank is 1.", @@ -111,11 +110,11 @@ class ContestWinner(BaseModel): description="When this user won this prize", ) - user: Optional[User] = Field(exclude=True, default=None) + user: User | None = Field(exclude=True, default=None) prize: ContestPrize = Field() - awarded_cash_amount: Optional[USDCent] = Field( + awarded_cash_amount: USDCent | None = Field( default=None, description="The actual amount this user receives. For cash prizes, if there was a tie, " "this could be different from the prize amount.", @@ -123,7 +122,7 @@ class ContestWinner(BaseModel): @computed_field() @property - def product_user_id(self) -> Optional[str]: + def product_user_id(self) -> str | None: # TODO: we'll have to pull username or censored emails or something if self.user: return self.user.product_user_id @@ -133,7 +132,7 @@ class ContestWinner(BaseModel): # def censored_product_user_id(self) -> str: # return censor_product_user_id(self.user) - def model_dump_mysql(self, contest_id: int) -> Dict[str, Any]: + def model_dump_mysql(self, contest_id: int) -> dict[str, Any]: data = self.model_dump(mode="json", exclude={"user"}) data["contest_id"] = contest_id diff --git a/generalresearch/models/thl/contest/contest.py b/generalresearch/models/thl/contest/contest.py index af5b136..6889dcb 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 datetime, timezone -from typing import Any, Dict, List, Optional, Tuple +from typing import Any from uuid import uuid4 from pydantic import ( @@ -41,7 +41,7 @@ class ContestBase(BaseModel, ABC): name: str = Field( max_length=128, description="Name of contest. Can be displayed to user." ) - description: Optional[str] = Field( + description: str | None = Field( default=None, max_length=2048, description="Description of contest. Can be displayed to user.", @@ -53,18 +53,18 @@ class ContestBase(BaseModel, ABC): end_condition: ContestEndCondition = Field() """Defines the conditions to win one or more prizes once the contest is ended""" - prizes: List[ContestPrize] = Field(default_factory=list, min_length=1) + prizes: list[ContestPrize] = Field(default_factory=list, min_length=1) starts_at: AwareDatetimeISO = Field( description="When the contest starts", default_factory=lambda: datetime.now(tz=timezone.utc), ) - terms_and_conditions: Optional[HttpUrl] = Field(default=None) + terms_and_conditions: HttpUrl | None = Field(default=None) status: ContestStatus = Field(default=ContestStatus.ACTIVE) - country_isos: Optional[CountryISOs] = Field( + country_isos: CountryISOs | None = Field( description="Contest is restricted to these countries. If null, all countries are allowed", default=None, ) @@ -80,7 +80,7 @@ class ContestBase(BaseModel, ABC): class Contest(ContestBase): - id: Optional[int] = Field( + id: int | None = Field( default=None, exclude=True, description="pk in db", @@ -100,17 +100,17 @@ class Contest(ContestBase): "entries being created/modified", ) - ended_at: Optional[AwareDatetimeISO] = Field( + ended_at: AwareDatetimeISO | None = Field( default=None, description="When the contest ended", ) - end_reason: Optional[ContestEndReason] = Field( + end_reason: ContestEndReason | None = Field( default=None, description="The reason the contest ended", ) - all_winners: Optional[List[ContestWinner]] = Field( + all_winners: list[ContestWinner] | None = Field( default=None, exclude=True, description="All prize winners of this contest", @@ -136,7 +136,7 @@ class Contest(ContestBase): # return True # return False - def should_end(self) -> Tuple[bool, Optional[ContestEndReason]]: + 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: @@ -145,12 +145,13 @@ class Contest(ContestBase): return False, None @abstractmethod - def select_winners(self) -> Optional[List[ContestWinner]]: ... + def select_winners(self) -> list[ContestWinner] | None: ... def end_contest(self) -> None: e, reason = self.should_end() if not e: - return None + return + # todo: Acquire a lock here, b/c this next part involves randomness # so we can't have it happen more than once winners = self.select_winners() @@ -169,7 +170,7 @@ class Contest(ContestBase): ) return None - def model_dump_mysql(self, **kwargs) -> Dict[str, Any]: + def model_dump_mysql(self, **kwargs) -> dict[str, Any]: d = self.model_dump(mode="json", **kwargs) d["created_at"] = self.created_at @@ -183,7 +184,7 @@ class Contest(ContestBase): return d @classmethod - def model_validate_mysql(cls, data: Dict[str, Any]) -> Self: + def model_validate_mysql(cls, data: dict[str, Any]) -> Self: data = {k: v for k, v in data.items() if k in cls.model_fields.keys()} if isinstance(data["end_condition"], dict): data["end_condition"] = ContestEndCondition.model_validate( @@ -204,12 +205,12 @@ class ContestUserView(Contest): # TODO: this could show a more detailed ContestWinner model, maybe # including like shipping status or whatever - user_winnings: List[ContestWinner] = Field( + user_winnings: list[ContestWinner] = Field( description="The prizes won in this contest by the requested user", default_factory=list, ) - def is_user_eligible(self, country_iso: str) -> Tuple[bool, str]: + def is_user_eligible(self, country_iso: str) -> tuple[bool, str]: now = datetime.now(tz=timezone.utc) assert country_iso.lower() == country_iso diff --git a/generalresearch/models/thl/contest/contest_entry.py b/generalresearch/models/thl/contest/contest_entry.py index 19586da..cddae14 100644 --- a/generalresearch/models/thl/contest/contest_entry.py +++ b/generalresearch/models/thl/contest/contest_entry.py @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import datetime, timezone -from typing import Any, Dict, Union from uuid import uuid4 from pydantic import ( @@ -20,7 +19,7 @@ from generalresearch.models.thl.user import User class ContestEntryCreate(BaseModel): entry_type: ContestEntryType = Field() # The meaning of this field is dictated by the contest's ContestEntryType - amount: Union[USDCent, int] = Field( + amount: USDCent | int = Field( description="The amount of the entry in integer counts or USD Cents", gt=0, ) @@ -51,7 +50,7 @@ class ContestEntry(BaseModel): entry_type: ContestEntryType = Field() # The meaning of this field is dictated by the contest's ContestEntryType - amount: Union[USDCent, int] = Field( + amount: USDCent | int = Field( description="The amount of the entry in integer counts or USD Cents", gt=0, ) @@ -61,7 +60,7 @@ class ContestEntry(BaseModel): @model_validator(mode="before") @classmethod - def validate_amount_type(cls, data: Dict) -> Dict: + def validate_amount_type(cls, data: dict) -> dict: from generalresearch.models.thl.contest.definitions import ( ContestEntryType, ) @@ -102,7 +101,7 @@ class ContestEntry(BaseModel): return censor_product_user_id(user=self.user) - def model_dump_mysql(self, contest_id: int) -> Dict[str, Any]: + def model_dump_mysql(self, contest_id: int) -> dict[str, Any]: data = self.model_dump(mode="json", exclude={"user"}) data["contest_id"] = contest_id data["created_at"] = self.created_at diff --git a/generalresearch/models/thl/contest/examples.py b/generalresearch/models/thl/contest/examples.py index 4835c14..1810e63 100644 --- a/generalresearch/models/thl/contest/examples.py +++ b/generalresearch/models/thl/contest/examples.py @@ -1,4 +1,6 @@ -from typing import Any, Dict +from __future__ import annotations + +from typing import Any from pydantic import HttpUrl @@ -6,7 +8,7 @@ from generalresearch.config import EXAMPLE_PRODUCT_ID from generalresearch.currency import USDCent -def _example_raffle_create(schema: Dict[str, Any]) -> None: +def _example_raffle_create(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestEndCondition, ContestEntryRule, @@ -46,7 +48,7 @@ def _example_raffle_create(schema: Dict[str, Any]) -> None: ).model_dump(mode="json") -def _example_raffle(schema: Dict) -> None: +def _example_raffle(schema: dict) -> None: from generalresearch.models.thl.contest import ( ContestEndCondition, ContestEntryRule, @@ -94,7 +96,7 @@ def _example_raffle(schema: Dict) -> None: return None -def _example_raffle_user_view(schema: Dict[str, Any]) -> None: +def _example_raffle_user_view(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestEndCondition, ContestEntryRule, @@ -145,7 +147,7 @@ def _example_raffle_user_view(schema: Dict[str, Any]) -> None: return None -def _example_milestone_create(schema: Dict[str, Any]) -> None: +def _example_milestone_create(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestPrize, ) @@ -186,7 +188,7 @@ def _example_milestone_create(schema: Dict[str, Any]) -> None: return None -def _example_milestone(schema: Dict[str, Any]) -> None: +def _example_milestone(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestPrize, ) @@ -232,7 +234,7 @@ def _example_milestone(schema: Dict[str, Any]) -> None: return None -def _example_milestone_user_view(schema: Dict[str, Any]) -> None: +def _example_milestone_user_view(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ContestPrize from generalresearch.models.thl.contest.definitions import ( ContestPrizeKind, @@ -278,7 +280,7 @@ def _example_milestone_user_view(schema: Dict[str, Any]) -> None: return None -def _example_leaderboard_contest_create(schema: Dict[str, Any]) -> None: +def _example_leaderboard_contest_create(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestPrize, ) @@ -323,7 +325,7 @@ def _example_leaderboard_contest_create(schema: Dict[str, Any]) -> None: return None -def _example_leaderboard_contest(schema: Dict[str, Any]) -> None: +def _example_leaderboard_contest(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestPrize, ) @@ -369,7 +371,7 @@ def _example_leaderboard_contest(schema: Dict[str, Any]) -> None: return None -def _example_leaderboard_contest_user_view(schema: Dict[str, Any]) -> None: +def _example_leaderboard_contest_user_view(schema: dict[str, Any]) -> None: from generalresearch.models.thl.contest import ( ContestPrize, ) @@ -412,5 +414,3 @@ def _example_leaderboard_contest_user_view(schema: Dict[str, Any]) -> None: product_id=EXAMPLE_PRODUCT_ID, product_user_id="test-user", ).model_dump(mode="json") - - return None diff --git a/generalresearch/models/thl/contest/io.py b/generalresearch/models/thl/contest/io.py index 8133c1a..e68f76e 100644 --- a/generalresearch/models/thl/contest/io.py +++ b/generalresearch/models/thl/contest/io.py @@ -1,5 +1,6 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Union from uuid import uuid4 from generalresearch.models.thl.contest.definitions import ContestType @@ -29,9 +30,7 @@ user_model_cls = { ContestType.MILESTONE: MilestoneUserView, ContestType.LEADERBOARD: LeaderboardContestUserView, } -ContestCreate = Union[ - RaffleContestCreate, LeaderboardContestCreate, MilestoneContestCreate -] +ContestCreate = RaffleContestCreate | LeaderboardContestCreate | MilestoneContestCreate from generalresearch.models.thl.contest.contest import Contest diff --git a/generalresearch/models/thl/contest/leaderboard.py b/generalresearch/models/thl/contest/leaderboard.py index d0681c4..8167f46 100644 --- a/generalresearch/models/thl/contest/leaderboard.py +++ b/generalresearch/models/thl/contest/leaderboard.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Literal, Optional, Tuple +from typing import Any, Literal from pydantic import ( ConfigDict, @@ -96,7 +98,7 @@ class LeaderboardContestCreate(ContestBase): return self @property - def leaderboard_key_parts(self) -> Dict[str, Any]: + def leaderboard_key_parts(self) -> dict[str, Any]: assert self.leaderboard_key.count(":") == 5, "invalid leaderboard_key" parts = self.leaderboard_key.split(":") _, product_id, country_iso, freq_str, date_str, board_code_value = parts @@ -137,8 +139,8 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): default=LeaderboardTieBreakStrategy.SPLIT_PRIZE_POOL ) - _redis_client: Optional[Redis] = PrivateAttr(default=None) - _user_manager: Optional[UserManager] = PrivateAttr(default=None) + _redis_client: Redis | None = PrivateAttr(default=None) + _user_manager: UserManager | None = PrivateAttr(default=None) @model_validator(mode="after") def validate_product_lb_key(self) -> Self: @@ -190,7 +192,7 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): ) return lbm - def should_end(self) -> Tuple[bool, Optional[ContestEndReason]]: + 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: @@ -198,7 +200,7 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): return False, None - def select_winners(self) -> List[ContestWinner]: + def select_winners(self) -> list[ContestWinner]: from generalresearch.models.thl.contest.utils import ( distribute_leaderboard_prizes, ) @@ -233,7 +235,7 @@ class LeaderboardContest(LeaderboardContestCreate, Contest): def country_iso(self) -> str: return self.leaderboard_key.split(":")[2] - def model_dump_mysql(self) -> Dict[str, Any]: + def model_dump_mysql(self) -> dict[str, Any]: d = super().model_dump_mysql( exclude={ "tie_break_strategy", @@ -252,7 +254,7 @@ class LeaderboardContestUserView(LeaderboardContest, ContestUserView): @computed_field(description="The current rank of this user in this contest") @property - def user_rank(self) -> Optional[int]: + def user_rank(self) -> int | None: if not self._redis_client: return None @@ -263,7 +265,7 @@ class LeaderboardContestUserView(LeaderboardContest, ContestUserView): return None - def is_user_eligible(self, country_iso: str) -> Tuple[bool, str]: + def is_user_eligible(self, country_iso: str) -> tuple[bool, str]: passes, msg = super().is_user_eligible(country_iso=country_iso) if not passes: return False, msg diff --git a/generalresearch/models/thl/contest/milestone.py b/generalresearch/models/thl/contest/milestone.py index 8c51a88..f62be8f 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, Dict, Literal, Optional, Tuple +from typing import Any, Literal from pydantic import ( BaseModel, @@ -53,13 +53,13 @@ class MilestoneContestEndCondition(BaseModel): Multiple conditions can be set. The contest is over once ANY conditions are met. """ - max_winners: Optional[PositiveInt] = Field( + max_winners: PositiveInt | None = Field( default=None, description="The contest will end once this many users have won (i.e. reached" "the milestone).", ) - ends_at: Optional[AwareDatetimeISO] = Field( + ends_at: AwareDatetimeISO | None = Field( default=None, description="The Contest is over at the ends_at time." ) @@ -72,17 +72,17 @@ class MilestoneContestConfig(BaseModel): target_amount: PositiveInt = Field( description="Each user 'wins' (receives prizes) once this target amount is reached." ) - entry_trigger: Optional[ContestEntryTrigger] = Field( + entry_trigger: ContestEntryTrigger | None = Field( description="What user action triggers an entry automatically.", default=None, ) # These two fields allow something like: "Get a complete in your first 24 hours!" - valid_for: Optional[timedelta] = Field( + valid_for: timedelta | None = Field( description="The time after valid_for_event for which the contest is open", default=None, ) - valid_for_event: Optional[Literal["signup"]] = Field(default=None) + valid_for_event: Literal["signup"] | None = Field(default=None) class MilestoneContestCreate(ContestBase, MilestoneContestConfig): @@ -127,7 +127,7 @@ class MilestoneContest(MilestoneContestCreate, Contest): default=0, ) - def should_end(self) -> Tuple[bool, Optional[ContestEndReason]]: + def should_end(self) -> tuple[bool, ContestEndReason | None]: res, msg = super().should_end() if res: @@ -145,7 +145,7 @@ class MilestoneContest(MilestoneContestCreate, Contest): # just does nothing return None - def model_dump_mysql(self) -> Dict[str, Any]: + def model_dump_mysql(self) -> dict[str, Any]: d = super().model_dump_mysql( exclude={ "entry_trigger", @@ -163,7 +163,7 @@ class MilestoneContest(MilestoneContestCreate, Contest): return d @classmethod - def model_validate_mysql(cls, data: Dict[str, Any]) -> Self: + def model_validate_mysql(cls, data: dict[str, Any]) -> Self: data.update( MilestoneContestConfig.model_validate(data["milestone_config"]).model_dump() ) @@ -180,7 +180,7 @@ class MilestoneUserView(MilestoneContest, ContestUserView): json_schema_extra=_example_milestone_user_view, ) - valid_until: Optional[AwareDatetimeISO] = Field( + valid_until: AwareDatetimeISO | None = Field( default=None, exclude=True, description="If valid_for is set, this gets populated wrt this user", @@ -201,7 +201,7 @@ class MilestoneUserView(MilestoneContest, ContestUserView): return True return False - def is_user_eligible(self, country_iso: str) -> Tuple[bool, str]: + def is_user_eligible(self, country_iso: str) -> tuple[bool, str]: passes, msg = super().is_user_eligible(country_iso=country_iso) if not passes: return False, msg diff --git a/generalresearch/models/thl/contest/raffle.py b/generalresearch/models/thl/contest/raffle.py index 1592857..4f868a0 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 datetime, timezone -from typing import Any, Dict, List, Literal, Optional, Tuple, Union +from typing import Any, Literal from pydantic import ( ConfigDict, @@ -72,9 +72,9 @@ class RaffleContest(RaffleContestCreate, Contest): json_schema_extra=_example_raffle, ) - entries: List[ContestEntry] = Field(default_factory=list, exclude=True) + entries: list[ContestEntry] = Field(default_factory=list, exclude=True) - current_amount: Union[int, USDCent] = Field( + current_amount: int | USDCent = Field( default=0, description="Sum of all entry amounts" ) current_participants: int = Field( @@ -82,7 +82,7 @@ class RaffleContest(RaffleContestCreate, Contest): ) @field_validator("entries", mode="after") - def sort_entries(cls, v: List[ContestEntry]): + def sort_entries(cls, v: list[ContestEntry]): return sorted(v, key=lambda x: x.created_at) @model_validator(mode="after") @@ -116,7 +116,7 @@ class RaffleContest(RaffleContestCreate, Contest): ) return self - def select_winners(self) -> List["ContestWinner"]: + def select_winners(self) -> list["ContestWinner"]: from generalresearch.models.thl.contest import ContestWinner assert self.is_complete(), "contest must be complete to select a winner" @@ -146,7 +146,7 @@ class RaffleContest(RaffleContestCreate, Contest): return winners - def should_end(self) -> Tuple[bool, Optional["ContestEndReason"]]: + def should_end(self) -> tuple[bool, "ContestEndReason" | None]: res, msg = super().should_end() if res: return res, msg @@ -157,7 +157,7 @@ class RaffleContest(RaffleContestCreate, Contest): return False, None @staticmethod - def select_winner(user_amount: Dict[int, int]) -> int: + def select_winner(user_amount: dict[int, int]) -> int: """ user_amount: Dict[user_id, amount], is total entry count for each user, e.g. {1111: 5, 2222: 1, 3333: 2} @@ -186,10 +186,10 @@ class RaffleContest(RaffleContestCreate, Contest): def get_current_participants(self) -> int: return len({entry.user.user_id for entry in self.entries}) - def get_current_amount(self) -> Union[int, USDCent]: + def get_current_amount(self) -> int | USDCent: return sum([x.amount for x in self.entries]) - def get_user_amount(self, product_user_id: str) -> Union[int, USDCent]: + def get_user_amount(self, product_user_id: str) -> int | USDCent: # Sum of this user's amounts return sum( e.amount for e in self.entries if e.user.product_user_id == product_user_id @@ -206,13 +206,13 @@ class RaffleContest(RaffleContestCreate, Contest): return True return False - def model_dump_mysql(self) -> Dict[str, Any]: + def model_dump_mysql(self) -> dict[str, Any]: d = super().model_dump_mysql() d["entry_rule"] = self.entry_rule.model_dump_json() return d @classmethod - def model_validate_mysql(cls, data: Dict) -> Self: + def model_validate_mysql(cls, data: dict) -> Self: data["entry_rule"] = ContestEntryRule.model_validate(data["entry_rule"]) return super().model_validate_mysql(data) @@ -224,10 +224,10 @@ class RaffleUserView(RaffleContest, ContestUserView): json_schema_extra=_example_raffle_user_view, ) - user_amount: Union[int, USDCent] = Field( + user_amount: int | USDCent = Field( description="The total amount this user has entered" ) - user_amount_today: Union[int, USDCent] = Field( + user_amount_today: int | USDCent = Field( description="The total amount this user has entered in the past 24 hours" ) @@ -252,7 +252,7 @@ class RaffleUserView(RaffleContest, ContestUserView): "end condition." ) @property - def projected_win_probability(self) -> Optional[float]: + def projected_win_probability(self) -> float | None: if self.end_condition.target_entry_amount is None: return None @@ -266,7 +266,7 @@ class RaffleUserView(RaffleContest, ContestUserView): # Not sure how to return this in api response, too confusing. Maybe use later. # Left for tests only. @property - def current_prize_count_probability(self) -> Dict[int, float]: + def current_prize_count_probability(self) -> int | float: # M: Population size (total entry amount) M = self.current_amount # n: number of success states (user's entry amount) @@ -278,7 +278,7 @@ class RaffleUserView(RaffleContest, ContestUserView): probs = {k: hypergeom.pmf(k, M, n, N) for k in range(1, N + 1)} return probs - def is_entry_eligible(self, entry: ContestEntry) -> Tuple[bool, str]: + def is_entry_eligible(self, entry: ContestEntry) -> tuple[bool, str]: if self.entry_rule.max_entry_amount_per_user: if ( self.user_amount + entry.amount @@ -292,7 +292,7 @@ class RaffleUserView(RaffleContest, ContestUserView): return False, "Entry would exceed max amount per user per day." return True, "" - def is_user_eligible(self, country_iso: str) -> Tuple[bool, str]: + def is_user_eligible(self, country_iso: str) -> tuple[bool, str]: passes, msg = super().is_user_eligible(country_iso=country_iso) if not passes: return False, msg diff --git a/generalresearch/models/thl/contest/utils.py b/generalresearch/models/thl/contest/utils.py index e5043c2..5ac2bc1 100644 --- a/generalresearch/models/thl/contest/utils.py +++ b/generalresearch/models/thl/contest/utils.py @@ -1,4 +1,6 @@ -from typing import TYPE_CHECKING, Dict, List +from __future__ import annotations + +from typing import TYPE_CHECKING if TYPE_CHECKING: from generalresearch.currency import USDCent @@ -18,8 +20,8 @@ def censor_product_user_id(user: "User") -> str: def distribute_leaderboard_prizes( - prizes: List["USDCent"], leaderboard_rows: List["LeaderboardRow"] -) -> Dict[str, "USDCent"]: + prizes: list["USDCent"], leaderboard_rows: list["LeaderboardRow"] +) -> dict[str, "USDCent"]: """ Distributes leaderboard prizes among tied users. The prizes for the tied places are pooled together and divided diff --git a/generalresearch/models/thl/demographics.py b/generalresearch/models/thl/demographics.py index a688e49..4d4c8c2 100644 --- a/generalresearch/models/thl/demographics.py +++ b/generalresearch/models/thl/demographics.py @@ -4,7 +4,7 @@ import copy from collections import Counter, defaultdict from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal +from typing import TYPE_CHECKING, Any, Literal import numpy as np @@ -76,7 +76,7 @@ class AgeGroup(Enum): return self.label -def calculate_demographic_metrics(opps: List[MarketplaceTask]) -> List: +def calculate_demographic_metrics(opps: list[MarketplaceTask]) -> list: """ Measurement: marketplace_survey_demographics tags: source (marketplace) @@ -146,8 +146,8 @@ def calculate_demographic_metrics(opps: List[MarketplaceTask]) -> List: def calculate_used_question_metrics( - opps: List[MarketplaceTask], qid_label: Dict[str, str] -) -> List[Dict[str, Any]]: + opps: list[MarketplaceTask], qid_label: dict[str, str] +) -> list[dict[str, Any]]: """ Measurement: marketplace_survey_targeting tags: source (marketplace), "type", country (all and individual) diff --git a/generalresearch/models/thl/finance.py b/generalresearch/models/thl/finance.py index 8d526ac..9ececc3 100644 --- a/generalresearch/models/thl/finance.py +++ b/generalresearch/models/thl/finance.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import random from datetime import timezone -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING from uuid import uuid4 import pandas as pd @@ -51,7 +53,7 @@ class POPFinancial(BaseModel): """ # --- Tracking / Tagging --- - product_id: Optional[UUIDStr] = Field(default=None, examples=[uuid4().hex]) + product_id: UUIDStr | None = Field(default=None, examples=[uuid4().hex]) time: AwareDatetimeISO = Field( description="The starting time block for the respective 'Period' that" @@ -77,7 +79,7 @@ class POPFinancial(BaseModel): examples=[adjustment_example], ) - adjustment_types: List[AdjustmentType] = Field() + adjustment_types: list[AdjustmentType] = Field() expense: int = Field( description="For Product accounts that are setup with Respondent payouts," @@ -101,8 +103,8 @@ class POPFinancial(BaseModel): @staticmethod def list_from_pandas( - input_data: pd.DataFrame, accounts: List["LedgerAccount"] - ) -> List["POPFinancial"]: + input_data: pd.DataFrame, accounts: list["LedgerAccount"] + ) -> list["POPFinancial"]: """ This list can either be for a Product or a Business. The difference is that the list of accounts will either be len()=1 (Product) or @@ -183,8 +185,8 @@ class ProductBalances(BaseModel): model_config = ConfigDict(extra="ignore", populate_by_name=True) # --- Tracking / Tagging --- - product_id: Optional[UUIDStr] = Field(default=None, examples=[uuid4().hex]) - last_event: Optional[AwareDatetimeISO] = Field(default=None) + product_id: UUIDStr | None = Field(default=None, examples=[uuid4().hex]) + last_event: AwareDatetimeISO | None = Field(default=None) # --- Numeric --- @@ -539,11 +541,11 @@ class ProductBalances(BaseModel): class BusinessBalances(BaseModel): - product_balances: List[ProductBalances] = Field(default_factory=list) + product_balances: list[ProductBalances] = Field(default_factory=list) # --- Validators --- @field_validator("product_balances") - def required_product_ids(cls, v: List[ProductBalances]): + def required_product_ids(cls, v: list[ProductBalances]): """The BusinessBalances needs to be able to distinguish between all the child Products; in order to do this, we need to assert that they all explicitly are set @@ -863,7 +865,7 @@ class BusinessBalances(BaseModel): from generalresearch.managers.thl.product import ProductManager pm = ProductManager(pg_config=thl_pg_config) - products: List[Product] = pm.get_by_uuids( + products: list[Product] = pm.get_by_uuids( product_uuids=[pb.product_id for pb in product_balances] ) sorted_products_uuids = [ diff --git a/generalresearch/models/thl/ipinfo.py b/generalresearch/models/thl/ipinfo.py index a98689f..5a8e51d 100644 --- a/generalresearch/models/thl/ipinfo.py +++ b/generalresearch/models/thl/ipinfo.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import ipaddress from datetime import datetime, timezone -from typing import Any, Dict, Literal, Optional, Tuple +from typing import Any, Literal import geoip2.models from faker import Faker @@ -27,7 +29,7 @@ fake = Faker() PrefixLength = Literal["/128", "/64", "/32"] -def normalize_ip(ip: IPvAnyAddressStr) -> Tuple[str, PrefixLength]: +def normalize_ip(ip: IPvAnyAddressStr) -> tuple[str, PrefixLength]: """ Normalize an IP address for MySQL storage. @@ -47,51 +49,51 @@ def normalize_ip(ip: IPvAnyAddressStr) -> Tuple[str, PrefixLength]: class IPGeoname(BaseModel): geoname_id: PositiveInt = Field() - continent_code: Optional[str] = Field(default=None, max_length=2) - continent_name: Optional[str] = Field(default=None, max_length=32) + continent_code: str | None = Field(default=None, max_length=2) + continent_name: str | None = Field(default=None, max_length=32) country_iso: CountryISOLike = Field( description="The ISO code of the country associated with the IP address.", examples=[fake.country_code().lower()], ) - country_name: Optional[str] = Field(default=None, max_length=64) + country_name: str | None = Field(default=None, max_length=64) - subdivision_1_iso: Optional[str] = Field( + subdivision_1_iso: str | None = Field( default=None, description="The ISO code of the primary subdivision (e.g., state or province).", max_length=3, ) - subdivision_1_name: Optional[str] = Field( + subdivision_1_name: str | None = Field( default=None, description="The name of the primary subdivision (e.g., state or province).", max_length=255, ) - subdivision_2_iso: Optional[str] = Field( + subdivision_2_iso: str | None = Field( default=None, description="The ISO code of the secondary subdivision (if applicable).", max_length=3, ) - subdivision_2_name: Optional[str] = Field( + subdivision_2_name: str | None = Field( default=None, description="The name of the secondary subdivision (if applicable).", max_length=255, ) - city_name: Optional[str] = Field( + city_name: str | None = Field( default=None, max_length=255, description="The name of the city associated with the IP address.", examples=[fake.city()], ) - metro_code: Optional[int] = Field(default=None) + metro_code: int | None = Field(default=None) - time_zone: Optional[str] = Field( + time_zone: str | None = Field( default=None, max_length=60, description="The time zone associated with the geographical location.", examples=[fake.timezone()], ) - is_in_european_union: Optional[bool] = Field(default=None) + is_in_european_union: bool | None = Field(default=None) updated: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc), @@ -104,20 +106,20 @@ class IPGeoname(BaseModel): "subdivision_2_iso", mode="before", ) - def make_lower(cls, value: Optional[str]) -> Optional[str]: + def make_lower(cls, value: str | None) -> str | None: if value is not None: return value.lower() return value # --- ORM --- - def model_dump_mysql(self) -> Dict[str, Any]: + def model_dump_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json") d["updated"] = self.updated return d @classmethod - def from_mysql(cls, d: Dict[str, Any]) -> Self: + def from_mysql(cls, d: dict[str, Any]) -> Self: d["updated"] = d["updated"].replace(tzinfo=timezone.utc) return cls.model_validate(d) @@ -167,77 +169,77 @@ class IPGeoname(BaseModel): class IPInformation(BaseModel): ip: IPvAnyAddressStr = Field() # This doesn't get stored in mysql/redis, b/c we only look up by the normalized ip - lookup_prefix: Optional[PrefixLength] = Field(default=None, exclude=True) + lookup_prefix: PrefixLength | None = Field(default=None, exclude=True) - geoname_id: Optional[PositiveInt] = Field(default=None) + geoname_id: PositiveInt | None = Field(default=None) country_iso: CountryISOLike = Field( description="The ISO code of the country associated with the IP address.", examples=[fake.country_code().lower()], ) - registered_country_iso: Optional[CountryISOLike] = Field( + registered_country_iso: CountryISOLike | None = Field( default=None, description="The ISO code of the country where the IP address is " "registered.", examples=[fake.country_code().lower()], ) - is_anonymous: Optional[bool] = Field( + is_anonymous: bool | None = Field( default=None, description="Indicates whether the IP address is associated with an " "anonymous source (e.g., VPN, proxy).", examples=[False], ) - is_anonymous_vpn: Optional[bool] = Field(default=None) - is_hosting_provider: Optional[bool] = Field(default=None) - is_public_proxy: Optional[bool] = Field(default=None) - is_tor_exit_node: Optional[bool] = Field(default=None) - is_residential_proxy: Optional[bool] = Field(default=None) + is_anonymous_vpn: bool | None = Field(default=None) + is_hosting_provider: bool | None = Field(default=None) + is_public_proxy: bool | None = Field(default=None) + is_tor_exit_node: bool | None = Field(default=None) + is_residential_proxy: bool | None = Field(default=None) - autonomous_system_number: Optional[PositiveInt] = Field(default=None) - autonomous_system_organization: Optional[str] = Field(default=None, max_length=255) + autonomous_system_number: PositiveInt | None = Field(default=None) + autonomous_system_organization: str | None = Field(default=None, max_length=255) - domain: Optional[str] = Field(default=None, max_length=255) - isp: Optional[str] = Field( + domain: str | None = Field(default=None, max_length=255) + isp: str | None = Field( default=None, description="The Internet Service Provider associated with the " "IP address.", examples=["Comcast"], ) - mobile_country_code: Optional[str] = Field(default=None, max_length=3) - mobile_network_code: Optional[str] = Field(default=None, max_length=3) + mobile_country_code: str | None = Field(default=None, max_length=3) + mobile_network_code: str | None = Field(default=None, max_length=3) - network: Optional[str] = Field(default=None, max_length=56) - organization: Optional[str] = Field(default=None, max_length=255) + network: str | None = Field(default=None, max_length=56) + organization: str | None = Field(default=None, max_length=255) - static_ip_score: Optional[float] = Field( + static_ip_score: float | None = Field( default=None, description="A score indicating the likelihood that the IP address is static.", ) - user_type: Optional[UserType] = Field( + user_type: UserType | None = Field( default=None, description="The type of user associated with the IP address " "(e.g., 'residential', 'business').", examples=[UserType.SCHOOL], ) - postal_code: Optional[str] = Field( + postal_code: str | None = Field( default=None, description="The postal code associated with the IP address.", examples=[fake.postcode()], ) - latitude: Optional[float] = Field( + latitude: float | None = Field( description="The latitude coordinate of the IP address location.", default=None, examples=[float(fake.latitude())], ) - longitude: Optional[float] = Field( + longitude: float | None = Field( description="The longitude coordinate of the IP address location.", default=None, examples=[float(fake.longitude())], ) - accuracy_radius: Optional[int] = Field( + accuracy_radius: int | None = Field( default=None, description="The approximate radius of accuracy for the latitude " "and longitude, in kilometers.", @@ -248,10 +250,10 @@ class IPInformation(BaseModel): default_factory=lambda: datetime.now(tz=timezone.utc), ) - _geoname: Optional[IPGeoname] = PrivateAttr(default=None) + _geoname: IPGeoname | None = PrivateAttr(default=None) @field_validator("country_iso", "registered_country_iso", mode="before") - def make_lower(cls, value: Optional[str]) -> Optional[str]: + def make_lower(cls, value: str | None) -> str | None: if value is not None: return value.lower() @@ -265,14 +267,13 @@ class IPInformation(BaseModel): return self.is_anonymous is None @property - def geoname(self) -> Optional["IPGeoname"]: + def geoname(self) -> "IPGeoname" | None: return self._geoname or None def normalize_ip(self): normalized_ip, lookup_prefix = normalize_ip(self.ip) self.ip = normalized_ip self.lookup_prefix = lookup_prefix - return None # --- prefetch_* --- def prefetch_geoname( @@ -288,8 +289,6 @@ class IPInformation(BaseModel): self._geoname = ip_gm.get_by_id(geoname_id=self.geoname_id) - return None - # --- ORM --- def model_dump_mysql(self): d = self.model_dump(mode="json", exclude={"geoname"}) @@ -297,7 +296,7 @@ class IPInformation(BaseModel): return d @classmethod - def from_mysql(cls, d: Dict) -> Self: + def from_mysql(cls, d: dict) -> Self: d["updated"] = d["updated"].replace(tzinfo=timezone.utc) return cls.model_validate(d) diff --git a/generalresearch/models/thl/leaderboard.py b/generalresearch/models/thl/leaderboard.py index 7d79091..399a906 100644 --- a/generalresearch/models/thl/leaderboard.py +++ b/generalresearch/models/thl/leaderboard.py @@ -4,7 +4,7 @@ import logging import math from datetime import datetime, timedelta, timezone from enum import Enum -from typing import List, Literal +from typing import Literal from uuid import UUID, uuid3 import pandas as pd @@ -117,7 +117,7 @@ class Leaderboard(BaseModel): description="The total number of rows in the leaderboard.", examples=[2] ) - rows: List[LeaderboardRow] = Field( + rows: list[LeaderboardRow] = Field( default_factory=list, examples=[ [ @@ -342,4 +342,4 @@ class LeaderboardWinner(BaseModel): class LeaderboardWinnerResponse(StatusResponse): - winners: List[LeaderboardWinner] = Field(default_factory=list) + winners: list[LeaderboardWinner] = Field(default_factory=list) diff --git a/generalresearch/models/thl/ledger.py b/generalresearch/models/thl/ledger.py index 9d473c5..3f8b123 100644 --- a/generalresearch/models/thl/ledger.py +++ b/generalresearch/models/thl/ledger.py @@ -1,6 +1,8 @@ +from __future__ import annotations + from datetime import datetime, timezone from enum import Enum -from typing import Annotated, Any, Dict, List, Literal, Optional, Union +from typing import Annotated, Any, Literal, Union from uuid import uuid4 from pydantic import ( @@ -164,9 +166,9 @@ class LedgerAccount(BaseModel, validate_assignment=True, frozen=True): normal_balance: Direction = Field(description=Direction.as_openapi()) - reference_type: Optional[str] = Field(default=None) + reference_type: str | None = Field(default=None) - reference_uuid: Optional[UUIDStr] = Field( + reference_uuid: UUIDStr | None = Field( default=None, description="The associated Product ID or other parent account that" "this Ledger Account is intended to track transactions for." @@ -202,7 +204,7 @@ class LedgerAccount(BaseModel, validate_assignment=True, frozen=True): class LedgerEntry(BaseModel): - id: Optional[int] = Field(default=None) + id: int | None = Field(default=None) direction: Direction account_uuid: UUIDStr @@ -217,7 +219,7 @@ class LedgerEntry(BaseModel): # This really shouldn't be Optional, but it has to be in order to # instantiate this class before the LedgerTransaction exists - transaction_id: Optional[int] = Field(default=None) + transaction_id: int | None = Field(default=None) @classmethod def from_amount(cls, account_uuid: UUIDStr, amount: int): @@ -248,7 +250,7 @@ class LedgerEntry(BaseModel): class LedgerTransaction(BaseModel): model_config = ConfigDict(extra="forbid", validate_assignment=True) - id: Optional[int] = Field(default=None) + id: int | None = Field(default=None) created: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc), @@ -258,11 +260,11 @@ class LedgerTransaction(BaseModel): "TX timestamps will likely be a few milliseconds delayed", ) - ext_description: Optional[str] = Field(default=None, max_length=255) - tag: Optional[str] = Field(default=None, max_length=255) - metadata: Dict[str, str] = Field(default_factory=dict) + ext_description: str | None = Field(default=None, max_length=255) + tag: str | None = Field(default=None, max_length=255) + metadata: dict[str, str] = Field(default_factory=dict) - entries: List[LedgerEntry] = Field( + entries: list[LedgerEntry] = Field( default_factory=list, description="A Transaction (TX) is composed of multiple Entry events.", ) @@ -288,7 +290,7 @@ class LedgerTransaction(BaseModel): @field_validator("entries", mode="after") @classmethod - def check_entries(cls, entries: List[LedgerEntry]) -> List[LedgerEntry]: + def check_entries(cls, entries: list[LedgerEntry]) -> list[LedgerEntry]: """Transactions should enforce double-entry upon creation. Each transaction needs to have at least two entries, which, in aggregate, must affect credit and debit sides in equal amounts. @@ -300,7 +302,7 @@ class LedgerTransaction(BaseModel): ), "ledger entries must balance" return entries - def model_dump_mysql(self, *args, **kwargs) -> Dict[str, Any]: + def model_dump_mysql(self, *args, **kwargs) -> dict[str, Any]: d = self.model_dump(mode="json", *args, **kwargs) if "created" in d: d["created"] = self.created.replace(tzinfo=None) @@ -387,12 +389,12 @@ class UserLedgerTransaction(BaseModel): ) # Needed to generate urls - product_id: Optional[str] = Field(default=None, exclude=True) + product_id: str | None = Field(default=None, exclude=True) # Needed to generate amount_string - payout_format: Optional[PayoutFormatType] = Field(default=None, exclude=True) + payout_format: PayoutFormatType | None = Field(default=None, exclude=True) # The balance in this account immediately after this tx. # It is optional b/c we'll calculate this from the query - balance_after: Optional[int] = Field(default=None) + balance_after: int | None = Field(default=None) def create_url(self, product_id: str): raise NotImplementedError() @@ -400,7 +402,7 @@ class UserLedgerTransaction(BaseModel): @computed_field( description="A link to where the user can get more details about this transaction", ) - def url(self) -> Optional[HttpsUrlStr]: + def url(self) -> HttpsUrlStr | None: if self.product_id is None: return None return self.create_url(product_id=self.product_id) @@ -408,7 +410,7 @@ class UserLedgerTransaction(BaseModel): @computed_field( description="The 'amount' with the payout_format applied.", ) - def amount_string(self) -> Optional[HttpsUrlStr]: + def amount_string(self) -> HttpsUrlStr | None: if self.payout_format is None: return None return format_payout_format( @@ -546,13 +548,13 @@ UserLedgerTransactionType = Annotated[ class UserLedgerTransactionTypeSummary(BaseModel): entry_count: NonNegativeInt = Field(default=0) - min_amount: Optional[int] = Field( + min_amount: int | None = Field( description="positive or negative USDCent", default=None ) - max_amount: Optional[int] = Field( + max_amount: int | None = Field( description="positive or negative USDCent", default=None ) - total_amount: Optional[int] = Field( + total_amount: int | None = Field( description="positive or negative USDCent", default=None ) @@ -578,7 +580,7 @@ class UserLedgerTransactions(Page): A (paginated) collection that holds transaction models that can be shown to a (wallet-managed) user. """ - transactions: List[UserLedgerTransactionType] = Field(default_factory=list) + transactions: list[UserLedgerTransactionType] = Field(default_factory=list) # The summary is w.r.t an optional time-filter. The transactions are # paginated so the counts won't necesarily match. In other words, the # summary is across all transaction in all pages, not this the transactions @@ -589,7 +591,7 @@ class UserLedgerTransactions(Page): def from_txs( cls, user_account: LedgerAccount, - txs: List[LedgerTransaction], + txs: list[LedgerTransaction], product_id: str, payout_format: str, summary: UserLedgerTransactionTypesSummary, @@ -617,11 +619,11 @@ class UserLedgerTransactions(Page): class LedgerAccountStatement(BaseModel): - id: Optional[int] = Field(default=None) + id: int | None = Field(default=None) account_uuid: UUIDStr - filter_str: Optional[str] = Field(default=None) + filter_str: str | None = Field(default=None) effective_at_lower_bound: AwareDatetimeISO effective_at_upper_bound: AwareDatetimeISO starting_balance: int = Field(lt=2**63 - 1, ge=0) ending_balance: int = Field(lt=2**63 - 1, ge=0) - sql_query: Optional[str] = Field(default=None) + sql_query: str | None = Field(default=None) diff --git a/generalresearch/models/thl/ledger_example.py b/generalresearch/models/thl/ledger_example.py index bf120cb..767be85 100644 --- a/generalresearch/models/thl/ledger_example.py +++ b/generalresearch/models/thl/ledger_example.py @@ -1,9 +1,11 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Any, Dict +from typing import Any from uuid import uuid4 -def _example_user_tx_payout(schema: Dict[str, Any]) -> None: +def _example_user_tx_payout(schema: dict[str, Any]) -> None: from generalresearch.models.thl.ledger import ( UserLedgerTransactionUserPayout, ) @@ -18,7 +20,7 @@ def _example_user_tx_payout(schema: Dict[str, Any]) -> None: ).model_dump(mode="json") -def _example_user_tx_bonus(schema: Dict[str, Any]) -> None: +def _example_user_tx_bonus(schema: dict[str, Any]) -> None: from generalresearch.models.thl.ledger import ( UserLedgerTransactionUserBonus, ) @@ -32,7 +34,7 @@ def _example_user_tx_bonus(schema: Dict[str, Any]) -> None: ).model_dump(mode="json") -def _example_user_tx_complete(schema: Dict[str, Any]) -> None: +def _example_user_tx_complete(schema: dict[str, Any]) -> None: from generalresearch.models.thl.ledger import ( UserLedgerTransactionTaskComplete, ) @@ -47,7 +49,7 @@ def _example_user_tx_complete(schema: Dict[str, Any]) -> None: ).model_dump(mode="json") -def _example_user_tx_adjustment(schema: Dict[str, Any]) -> None: +def _example_user_tx_adjustment(schema: dict[str, Any]) -> None: from generalresearch.models.thl.ledger import ( UserLedgerTransactionTaskAdjustment, ) diff --git a/generalresearch/models/thl/locales.py b/generalresearch/models/thl/locales.py index bff774a..85f0c32 100644 --- a/generalresearch/models/thl/locales.py +++ b/generalresearch/models/thl/locales.py @@ -1,4 +1,6 @@ -from typing import Annotated, Set +from __future__ import annotations + +from typing import Annotated from pydantic import AfterValidator @@ -9,8 +11,8 @@ from generalresearch.models.custom_types import ( ) locale_helper = Localelator() -COUNTRY_ISOS: Set[str] = locale_helper.get_all_countries() -LANGUAGE_ISOS: Set[str] = locale_helper.get_all_languages() +COUNTRY_ISOS: set[str] = locale_helper.get_all_countries() +LANGUAGE_ISOS: set[str] = locale_helper.get_all_languages() def is_valid_country_iso(v: str) -> str: @@ -28,5 +30,5 @@ CountryISO = Annotated[str, AfterValidator(is_valid_country_iso)] # 3-char ISO 639-2/B, lowercase LanguageISO = Annotated[str, AfterValidator(is_valid_language_iso)] -CountryISOs = Annotated[Set[CountryISO], to_comma_sep_str, from_comma_sep_str] -LanguageISOs = Annotated[Set[LanguageISO], to_comma_sep_str, from_comma_sep_str] +CountryISOs = Annotated[set[CountryISO], to_comma_sep_str, from_comma_sep_str] +LanguageISOs = Annotated[set[LanguageISO], to_comma_sep_str, from_comma_sep_str] diff --git a/generalresearch/models/thl/maxmind/__init__.py b/generalresearch/models/thl/maxmind/__init__.py deleted file mode 100644 index e69de29..0000000 --- a/generalresearch/models/thl/maxmind/__init__.py +++ /dev/null diff --git a/generalresearch/models/thl/maxmind/definitions.py b/generalresearch/models/thl/maxmind/definitions.py deleted file mode 100644 index 01431c7..0000000 --- a/generalresearch/models/thl/maxmind/definitions.py +++ /dev/null @@ -1,22 +0,0 @@ -from enum import Enum - -from generalresearch.utils.enum import ReprEnumMeta - - -class UserType(Enum, metaclass=ReprEnumMeta): - # https://support.maxmind.com/hc/en-us/articles/4408430082971-IP-Trait-Risk-Data#h_01FN6V8JMQMWZGWNPPAW77ZPY4 - BUSINESS = "business" - CAFE = "cafe" - CELLULAR = "cellular" - COLLEGE = "college" - CDN = "content_delivery_network" - CPN = "consumer_privacy_network" - GOVERNMENT = "government" - HOSTING = "hosting" - LIBRARY = "library" - MILITARY = "military" - RESIDENTIAL = "residential" - ROUTER = "router" - SCHOOL = "school" - SEARCH_ENGINE = "search_engine_spider" - TRAVELER = "traveler" diff --git a/generalresearch/models/thl/offerwall/__init__.py b/generalresearch/models/thl/offerwall/__init__.py index bf6b7ea..d2d7d36 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, Dict, Literal, Optional, Set +from typing import Any, Literal from pydantic import ( BaseModel, @@ -111,7 +111,7 @@ class OfferWallRequest(BaseModel): offerwall_type: OfferWallType = Field() user: User = Field() - ip: Optional[IPvAnyAddressStr] = Field( + ip: IPvAnyAddressStr | None = Field( default=None, description="Respondent's IP address (IPv4 or IPv6). Either 'ip' must be " "provided, or 'country_iso' must be provided if 'ip' is " @@ -121,38 +121,38 @@ class OfferWallRequest(BaseModel): country_iso: CountryISO = Field( description="Respondent's country code (ISO 3166-1 alpha-2, lowercase)" ) - language_isos: Set[LanguageISO] = Field( + language_isos: set[LanguageISO] = Field( description="Respondent's desired language (ISO 639-2/B, lowercase)", ) - behavior: Optional[OfferWallBehaviorsType] = Field( + behavior: OfferWallBehaviorsType | None = Field( default=None, max_length=12, description="Allows using custom scoring functions. Please " "discuss directly with GRL.", ) - min_payout: Optional[Decimal] = Field( + min_payout: Decimal | None = Field( default=None, description="Decimal representation of the minimum amount of USD that " "any of the tasks will pay", examples=["1.23"], ) - duration: Optional[int] = Field( + duration: int | None = Field( default=60 * 90, description="Maximum length of desired task (in seconds).", gt=0, ) - n_bins: Optional[int] = Field( + n_bins: int | None = Field( default=None, description="Number of bins requested in the offerwall.", le=100, gt=0, ) - min_bin_size: Optional[int] = Field( + min_bin_size: int | None = Field( default=None, description="Minimum number of tasks that need to be in a bucket", gt=0, @@ -169,7 +169,7 @@ class OfferWallRequest(BaseModel): default="payout", description="Cluster tasks by payout or duration" ) - passthrough_kwargs: Dict[str, str] = Field( + passthrough_kwargs: dict[str, str] = Field( default_factory=dict, description="These are pulled from the url params. They are any 'extra' url params " "in the getofferwall request. They'll be available through the task_status " @@ -177,12 +177,12 @@ class OfferWallRequest(BaseModel): ) # Only for soft pair (offerwall_id, max_options, max_questions) - offerwall_id: Optional[str] = Field(default=None) - max_options: Optional[int] = Field( + offerwall_id: str | None = Field(default=None) + max_options: int | None = Field( default=None, description="Max number of options an allowed question can have (allowed to be asked)", ) - max_questions: Optional[int] = Field( + max_questions: int | None = Field( default=None, description="Max number of missing questions on a single bin", ) @@ -199,7 +199,7 @@ class OfferWallRequest(BaseModel): "these may be set in the bpc table globally for a BP.", ) - marketplaces: Optional[Set[Source]] = Field( + marketplaces: set[Source] | None = Field( default=None, description="If set, restrict tasks to those from these marketplaces only.", ) @@ -264,7 +264,7 @@ class OfferWallRequest(BaseModel): json.dumps(self.model_dump(mode="json"), sort_keys=True).encode("utf-8") ).hexdigest()[:7] - def to_grpc_request(self) -> Dict[str, Any]: + def to_grpc_request(self) -> dict[str, Any]: # We need this so thl-core can refresh an offerwall in order to continue # a session d = self.model_dump(mode="json") @@ -301,15 +301,15 @@ class OfferWallRequest(BaseModel): } @property - def product_id(self) -> Optional[str]: + def product_id(self) -> str | None: return self.user.product_id @property - def product_user_id(self) -> Optional[str]: + def product_user_id(self) -> str | None: return self.user.product_user_id @property - def bpuid(self) -> Optional[str]: + def bpuid(self) -> str | None: return self.user.product_user_id @property diff --git a/generalresearch/models/thl/offerwall/base.py b/generalresearch/models/thl/offerwall/base.py index 8ea259d..8e24958 100644 --- a/generalresearch/models/thl/offerwall/base.py +++ b/generalresearch/models/thl/offerwall/base.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import statistics from datetime import timedelta from decimal import Decimal from string import Formatter -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any from uuid import uuid4 import numpy as np @@ -218,7 +220,7 @@ class MergeTableFeatures(BaseModel): @model_validator(mode="before") @classmethod - def set_completion_time_log(cls, data: Dict[str, Any]) -> Dict[str, Any]: + def set_completion_time_log(cls, data: dict[str, Any]) -> dict[str, Any]: # This isn't actually in the merge table data["COMPLETION_TIME_LOG"] = np.log(data["COMPLETION_TIME"]) return data @@ -238,7 +240,7 @@ class TaskResult(BaseModel): ) source: Source = Field() country_iso: CountryISO = Field() - buyer_id: Optional[str] = Field(min_length=1, max_length=32, default=None) + buyer_id: str | None = Field(min_length=1, max_length=32, default=None) # todo: GRS is allowed to be 0, but all the others can't. make a validator cpi: Decimal = Field(ge=0, le=100, decimal_places=5, max_digits=7) @@ -246,7 +248,7 @@ class TaskResult(BaseModel): # Only GRS tasks will have this set. All other marketplaces will have # to make a grpc call to generate this. This is a str b/c it is actually # a format string. - entry_link: Optional[str] = Field( + entry_link: str | None = Field( default=None, examples=[ "https://{domain}/session/?39057c8b=c4ed212601494f8c8836e38a55102d10&c184efc0=test&0bb50182={mid}" @@ -327,12 +329,12 @@ class ScoredTaskResult(TaskResult, MergeTableFeatures): # The set of marketplace's question codes (internal id) that are unknown. # This should only be set it SoftPairResultType is conditional - unknown_mp_question_ids: Optional[Set[str]] = Field(default=None) + unknown_mp_question_ids: set[str] | None = Field(default=None) # Question ids (from marketplace_question table) for the questions that # will be asked (that would fulfill the unknown questions specified in # unknown_mp_question_ids) - unknown_question_ids: Optional[Set[UUIDStr]] = Field(default=None) + unknown_question_ids: set[UUIDStr] | None = Field(default=None) # ---- Soft Pair end ---- @@ -343,7 +345,7 @@ class ScoredTaskResult(TaskResult, MergeTableFeatures): return Decimal(v).quantize(Decimal("0.00000")) @property - def unknown_mp_qids(self) -> Optional[Set[str]]: + def unknown_mp_qids(self) -> set[str] | None: # marketplace's curie-formatted question IDs that are unknown return ( {self.source + ":" + q for q in self.unknown_mp_question_ids} @@ -351,14 +353,14 @@ class ScoredTaskResult(TaskResult, MergeTableFeatures): else None ) - def to_row(self) -> Dict[str, Any]: + def to_row(self) -> dict[str, Any]: d = self.model_dump(mode="json") d["id_code"] = self.id_code return d class ScoredTaskResults(BaseModel): - tasks: List[ScoredTaskResult] = Field() + tasks: list[ScoredTaskResult] = Field() @property def availability_count(self) -> NonNegativeInt: @@ -375,7 +377,7 @@ class ScoredTaskResults(BaseModel): df["cpi"] = df["cpi"].astype(float) return df - def take_top(self, n=100) -> List[ScoredTaskResult]: + def take_top(self, n=100) -> list[ScoredTaskResult]: return sorted(self.tasks, key=lambda x: x.score, reverse=True)[:n] @@ -395,7 +397,7 @@ class OfferwallBucket(BaseModel): examples=["5ba2fe5010cc4d078fc3cc0b0cc264c3"], default_factory=lambda: uuid4().hex, ) - uri: Optional[HttpsUrl] = Field( + uri: HttpsUrl | None = Field( examples=[ "https://task.generalresearch.com/api/v1/52d3f63b2709/797df4136c604a6c8599818296aae6d1/?i" "=5ba2fe5010cc4d078fc3cc0b0cc264c3&b=test&66482fb=e7baf5e" @@ -404,13 +406,13 @@ class OfferwallBucket(BaseModel): default=None, ) - tasks: List[ScoredTaskResult] = Field() + tasks: list[ScoredTaskResult] = Field() - category: List[CategoryAssociation] = Field(default_factory=list) + category: list[CategoryAssociation] = Field(default_factory=list) # Used only in marketplace offerwall - source: Optional[Source] = Field(default=None) - source_name: Optional[str] = Field(default=None) + source: Source | None = Field(default=None) + source_name: str | None = Field(default=None) # Normally these are calculated. However, in some offerwalls we duplicate # buckets, so they're not "true" calculated values. @@ -420,7 +422,7 @@ class OfferwallBucket(BaseModel): description="Custom: Min payout across all tasks", default=None, ) - custom_q1_duration: Optional[float] = Field( + custom_q1_duration: float | None = Field( description="Custom: Q1 loi across all tasks", default=None, gt=0, @@ -429,11 +431,11 @@ class OfferwallBucket(BaseModel): quality_score: float = Field(default=0) - eligibility_criteria: Optional[Tuple[SurveyEligibilityCriterion, ...]] = Field( + eligibility_criteria: tuple[SurveyEligibilityCriterion, ...] | None = Field( description="The reasons the user is eligible for tasks in this bucket", default=None, ) - eligibility_explanation: Optional[str] = Field( + eligibility_explanation: str | None = Field( default=None, description="Human-readable text explaining a user's eligibility for tasks in this bucket", examples=[ @@ -442,7 +444,7 @@ class OfferwallBucket(BaseModel): ) @property - def missing_questions(self) -> Set[UUIDStr]: + def missing_questions(self) -> set[UUIDStr]: # Used only in softpair. # The question id is the question's uuid (in the marketplace_question table / UpkQuestion.id) # It is just the set union of task.softpair.question_ids for all tasks in this bucket. @@ -461,12 +463,12 @@ class OfferwallBucket(BaseModel): return float(np.sqrt((score**2).sum())) @property - def payout(self) -> Optional[Decimal]: + def payout(self) -> Decimal | None: # The payout is the Min payout across all tasks return min([x.payout for x in self.tasks], default=None) @property - def loi(self) -> Optional[float]: + def loi(self) -> float | None: # The loi is the Max LOI across all tasks return max([x.loi for x in self.tasks], default=None) @@ -487,7 +489,7 @@ class OfferwallBucket(BaseModel): return self.duration_summary.max @property - def sns(self) -> List[str]: + def sns(self) -> list[str]: return [t.id_code for t in self.tasks] @property @@ -541,7 +543,7 @@ class OfferwallBucket(BaseModel): ) @property - def eligibility(self) -> Optional[SoftPairResultType]: + def eligibility(self) -> SoftPairResultType | None: # We're assuming there is never a conditional or ineligible survey # after a unconditional. There can be unconditional surveys # after conditional surveys, in which case the bucket is still @@ -558,7 +560,7 @@ class OfferwallBucket(BaseModel): raise ValueError(f"Unexpected pair_type {pair_type}") @property - def eligibility_str(self) -> Optional[Eligibility]: + def eligibility_str(self) -> Eligibility | None: return ( { SoftPairResultType.UNCONDITIONAL: "unconditional", @@ -586,7 +588,7 @@ class OfferwallBucket(BaseModel): ) def generate_bucket_entry_url( - self, user: User, request_id: Optional[str] = None + self, user: User, request_id: str | None = None ) -> None: product_id = user.product_id product_user_id = user.product_user_id @@ -601,8 +603,6 @@ class OfferwallBucket(BaseModel): base_enter_url, self.id, product_user_id, request_id=request_id ) - return None - # def __repr__(self): # exclude = { # "PRESCREEN_CONVERSION_ALPHA", @@ -646,7 +646,7 @@ class OfferwallBase(BaseModel): default_factory=lambda: uuid4().hex, ) offerwall_type: OfferWallType = Field() - buckets: List[OfferwallBucket] = Field() + buckets: list[OfferwallBucket] = Field() # Note: this != the sum(len(tasks) in buckets) b/c we filter out a lot availability_count: int = Field(default=0, description="Number of available tasks") @@ -660,7 +660,7 @@ class OfferwallBase(BaseModel): examples=[7], default=0, ) - offerwall_reasons: List[OfferwallReason] = Field( + offerwall_reasons: list[OfferwallReason] = Field( default_factory=list, description=( "Explanations describing why so many or few opportunities are available." @@ -670,7 +670,7 @@ class OfferwallBase(BaseModel): # Contains the full info about any questions in any bucket's # missing_questions. - questions: List[UpkQuestion] = Field(default_factory=list) + questions: list[UpkQuestion] = Field(default_factory=list) @property def offerwall_type_class(self) -> OfferWallTypeClass: @@ -683,5 +683,3 @@ class OfferwallBase(BaseModel): def generate_bucket_entry_urls(self, user: User, request_id: str) -> None: for bucket in self.buckets: bucket.generate_bucket_entry_url(user=user, request_id=request_id) - - return None diff --git a/generalresearch/models/thl/offerwall/behavior.py b/generalresearch/models/thl/offerwall/behavior.py index 8e6c089..555870d 100644 --- a/generalresearch/models/thl/offerwall/behavior.py +++ b/generalresearch/models/thl/offerwall/behavior.py @@ -1,4 +1,6 @@ -from typing import Any, Dict, Literal +from __future__ import annotations + +from typing import Any, Literal from pydantic import BaseModel, Field @@ -6,7 +8,7 @@ from pydantic import BaseModel, Field class OfferWallBehavior(BaseModel): id: str = Field() name: str = Field() - kwargs: Dict[str, Any] = Field(default_factory=dict) + kwargs: dict[str, Any] = Field(default_factory=dict) OFFERWALL_BEHAVIOR_PRESETS = [ diff --git a/generalresearch/models/thl/offerwall/bucket.py b/generalresearch/models/thl/offerwall/bucket.py index 1e08e32..bb18ebe 100644 --- a/generalresearch/models/thl/offerwall/bucket.py +++ b/generalresearch/models/thl/offerwall/bucket.py @@ -1,4 +1,5 @@ -from typing import Optional +from __future__ import annotations + from urllib.parse import urlencode @@ -6,8 +7,8 @@ def generate_offerwall_entry_url( base_url: str, obj_id: str, bp_user_id: str, - request_id: Optional[str] = None, - nudge_id: Optional[str] = None, + request_id: str | None = None, + nudge_id: str | None = None, ) -> str: # For an offerwall entry link, we need the clicked bucket_id and the # request hash (so we know which GetOfferwall cache to get diff --git a/generalresearch/models/thl/offerwall/cache.py b/generalresearch/models/thl/offerwall/cache.py index b75a803..c36568e 100644 --- a/generalresearch/models/thl/offerwall/cache.py +++ b/generalresearch/models/thl/offerwall/cache.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any from pydantic import BaseModel, Field @@ -23,21 +25,21 @@ class GetOfferWallCache(BaseModel): request: OfferWallRequest = Field() request_id: str = Field() offerwall: OfferwallBase = Field() - all_sids: List[str] = Field() + all_sids: list[str] = Field() timestamp: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(timezone.utc) ) - latest_ip_info: Dict[str, Any] = Field( + latest_ip_info: dict[str, Any] = Field( description="So we can easily check if user's IP info has changed" ) - profiling_task: Optional[TaskResult] = Field( + profiling_task: TaskResult | None = Field( description="Profiling task", default=None ) is_avg_offerwall: bool = Field() # These only get set once a bucket is clicked. - clicked_timestamp: Optional[AwareDatetimeISO] = Field(default=None) - clicked_bucket: Optional[UUIDStr] = Field(default=None) + clicked_timestamp: AwareDatetimeISO | None = Field(default=None) + clicked_bucket: UUIDStr | None = Field(default=None) class SessionInfoCache(BaseModel): @@ -47,13 +49,13 @@ class SessionInfoCache(BaseModel): # This starts out as just the tasks within the clicked bucket, but # will get pruned as tasks are attempted - tasks: List[ScoredTaskResult] = Field() + tasks: list[ScoredTaskResult] = Field() started: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc) ) # The count of attempts per marketplace - mp_retry_count: Dict[Source, int] = Field(default_factory=dict) + mp_retry_count: dict[Source, int] = Field(default_factory=dict) hard_retry_count: int = Field(default=0) diff --git a/generalresearch/models/thl/pagination.py b/generalresearch/models/thl/pagination.py index 1b31078..e8311f0 100644 --- a/generalresearch/models/thl/pagination.py +++ b/generalresearch/models/thl/pagination.py @@ -1,5 +1,6 @@ +from __future__ import annotations + from math import ceil -from typing import Optional from pydantic import BaseModel, Field, computed_field @@ -8,12 +9,10 @@ class Page(BaseModel): # Based on fastapi_pagination.Page page: int = Field(default=1, ge=1, description="Page number") size: int = Field(default=50, ge=1, le=100, description="Page size") - total: Optional[int] = Field( - default=None, ge=0, description="Total number of results" - ) + total: int | None = Field(default=None, ge=0, description="Total number of results") @computed_field(description="Total number of pages") - def pages(self) -> Optional[int]: + def pages(self) -> int | None: if self.size == 0: return 0 elif self.total is not None: diff --git a/generalresearch/models/thl/payout.py b/generalresearch/models/thl/payout.py index 8ab01ec..1a9d534 100644 --- a/generalresearch/models/thl/payout.py +++ b/generalresearch/models/thl/payout.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import json from datetime import datetime, timezone -from typing import Collection, Dict, List, Optional +from typing import Collection from uuid import uuid4 from pydantic import ( @@ -70,25 +72,25 @@ class PayoutEvent(BaseModel): examples=[531], ) - status: Optional[PayoutStatus] = Field( + status: PayoutStatus | None = Field( default=PayoutStatus.PENDING, description=PayoutStatus.as_openapi(), examples=[PayoutStatus.COMPLETE], ) # Used for holding an external, payout-type-specific identifier - ext_ref_id: Optional[str] = Field(default=None) + ext_ref_id: str | None = Field(default=None) payout_type: PayoutType = Field( description=PayoutType.as_openapi(), examples=[PayoutType.ACH] ) - request_data: Dict = Field( + request_data: dict = Field( default_factory=dict, description="Stores payout-type-specific information that is used to " "request this payout from the external provider.", ) - order_data: Optional[Dict | CashMailOrderData] = Field( + order_data: dict | CashMailOrderData | None = Field( default=None, description="Stores payout-type-specific order information that is " "returned from the external payout provider.", @@ -97,8 +99,8 @@ class PayoutEvent(BaseModel): def update( self, status: PayoutStatus, - ext_ref_id: Optional[str] = None, - order_data: Optional[Dict] = None, + ext_ref_id: str | None = None, + order_data: dict | None = None, ) -> None: self.check_status_change_allowed(status) @@ -108,15 +110,13 @@ class PayoutEvent(BaseModel): self.ext_ref_id = ext_ref_id self.order_data = order_data - return None - def check_status_change_allowed(self, status: PayoutStatus) -> None: # We may not be changing the status when this method gets called. It's # possible to be called when we're updating other attributes so # allow immediate bypass if it isn't actually different. if self.status == status: - return None + return if self.status in { PayoutStatus.REJECTED, @@ -144,7 +144,7 @@ class PayoutEvent(BaseModel): # --- ORM --- - def model_dump_mysql(self, *args, **kwargs) -> Dict: + def model_dump_mysql(self, *args, **kwargs) -> dict: d = self.model_dump(mode="json", *args, **kwargs) if "created" in d: @@ -171,13 +171,13 @@ class UserPayoutEvent(PayoutEvent): # Requires joining on: # - accounting_cashoutmethod # - ledger_account - account_reference_type: Optional[str] = Field(default=None) - account_reference_uuid: Optional[UUIDStr] = Field(default=None) + account_reference_type: str | None = Field(default=None) + account_reference_uuid: UUIDStr | None = Field(default=None) # By default, this will just be the cashout_method.name. This also is # populated from the db and so does not need to be set (there is no # `description` field in event_payout) - description: Optional[str] = Field(default=None) + description: str | None = Field(default=None) @field_validator("payout_type", mode="before") @classmethod @@ -226,14 +226,14 @@ class BrokerageProductPayoutEvent(PayoutEvent): def from_payout_event( cls, pe: PayoutEvent, - account_product_mapping: Optional[Dict[UUIDStr, UUIDStr]] = None, - redis_config: Optional[RedisConfig] = None, + account_product_mapping: dict[UUIDStr, UUIDStr] | None = None, + redis_config: RedisConfig | None = None, ) -> Self: # TODO!: prevent re-assignment, rework this... if account_product_mapping is None: rc = redis_config.create_redis_client() - account_product_mapping: Dict = rc.hgetall(name="pem:account_to_product") + account_product_mapping: dict = rc.hgetall(name="pem:account_to_product") assert isinstance(account_product_mapping, dict) assert pe.uuid in account_product_mapping.keys() @@ -246,14 +246,14 @@ class BrokerageProductPayoutEvent(PayoutEvent): cls, payout_events: Collection[PayoutEvent], order_by=OrderBy, - account_product_mapping: Optional[Dict[UUIDStr, UUIDStr]] = None, - redis_config: Optional[RedisConfig] = None, - ) -> List[Self]: + account_product_mapping: dict[UUIDStr, UUIDStr] | None = None, + redis_config: RedisConfig | None = None, + ) -> list[Self]: # TODO!: prevent re-assignment, rework this... if account_product_mapping is None: rc = redis_config.create_redis_client() - account_product_mapping: Dict = rc.hgetall(name="pem:account_to_product") + account_product_mapping: dict = rc.hgetall(name="pem:account_to_product") assert isinstance(account_product_mapping, dict) res = [] @@ -278,7 +278,7 @@ class BrokerageProductPayoutEvent(PayoutEvent): class BusinessPayoutEvent(BaseModel): """A single ACH or Wire event to a Business Bank Account""" - bp_payouts: List[BrokerageProductPayoutEvent] = Field( + bp_payouts: list[BrokerageProductPayoutEvent] = Field( description="Here is the list of Brokerage Product Payouts that" "this Business Payout includes.", min_length=1, @@ -326,7 +326,7 @@ class BusinessPayoutEvent(BaseModel): @computed_field( title="External Reference ID", description="ACH Transaction ID", - return_type=Optional[str], + return_type=str | None, ) @property def ext_ref_id(self): diff --git a/generalresearch/models/thl/payout_format.py b/generalresearch/models/thl/payout_format.py index 4989bdc..9f22ace 100644 --- a/generalresearch/models/thl/payout_format.py +++ b/generalresearch/models/thl/payout_format.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import decimal import re diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index eb542e1..47e50d0 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -13,11 +13,7 @@ from typing import ( TYPE_CHECKING, Any, Callable, - Dict, - List, Literal, - Optional, - Set, ) from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit from uuid import uuid4 @@ -138,7 +134,7 @@ class ProfilingConfig(BaseModel): the GRS system but won't present them any questions.""", ) - n_questions: Optional[PositiveInt] = Field( + n_questions: PositiveInt | None = Field( default=None, description="Use to hard code the number of questions to ask. None means use default algorithm.", ) @@ -187,7 +183,7 @@ class UserHealthConfig(BaseModel): # Users in these countries are "blocked". Blocked in quotes because # the user doesn't actually get blocked, they just are treated like they # are blocked. - banned_countries: List[CountryISOLike] = Field(default_factory=list) + banned_countries: list[CountryISOLike] = Field(default_factory=list) # Decide if a user can be blocked for IP-related triggers such as sharing # IPs and location history. This should eventually be deprecated and @@ -196,13 +192,13 @@ class UserHealthConfig(BaseModel): # These are only checked by ym-user-predict, which I'm not sure even # works properly. To be deprecated ... don't even use them. - userprofit_cutoff: Optional[Decimal] = Field(default=None, exclude=True) - recon_cutoff: Optional[float] = Field(default=None, exclude=True) - droprate_cutoff: Optional[float] = Field(default=None, exclude=True) - conversion_cutoff: Optional[float] = Field(default=None, exclude=True) + userprofit_cutoff: Decimal | None = Field(default=None, exclude=True) + recon_cutoff: float | None = Field(default=None, exclude=True) + droprate_cutoff: float | None = Field(default=None, exclude=True) + conversion_cutoff: float | None = Field(default=None, exclude=True) @field_validator("banned_countries", mode="after") - def sort_values(cls, values: List[str]): + def sort_values(cls, values: list[str]): return sorted(values) @@ -246,15 +242,15 @@ class OfferWallRequestYieldmanParams(BaseModel): class OfferWallCategoryRequest(BaseModel): # Only include these categories - adwords_category: Optional[List[str]] = Field(default=None, examples=[["45", "65"]]) - category: Optional[List[str]] = Field( + adwords_category: list[str] | None = Field(default=None, examples=[["45", "65"]]) + category: list[str] | None = Field( default=None, examples=[["98c137e4e90a4d92ac6c00e523eb1b50"]] ) # Exclude these categories - exclude_adwords_category: Optional[List[str]] = Field( + exclude_adwords_category: list[str] | None = Field( default=None, examples=[["1558"]] ) - exclude_category: Optional[List[str]] = Field( + exclude_category: list[str] | None = Field( default=None, examples=[ [ @@ -291,7 +287,7 @@ class SourcesConfig(BaseModel): model_config = ConfigDict(frozen=True) - user_defined: List[SourceConfig] = Field(default_factory=list) + user_defined: list[SourceConfig] = Field(default_factory=list) @model_validator(mode="after") def validate_user_defined(self): @@ -300,11 +296,11 @@ class SourcesConfig(BaseModel): return self @cached_property - def default_sources(self) -> List[SourceConfig]: + def default_sources(self) -> list[SourceConfig]: return [SourceConfig.model_validate({"name": s}) for s in Source] @cached_property - def sources(self) -> List[SourceConfig]: + def sources(self) -> list[SourceConfig]: # If a BP has no user_defined SourceConfigs, we use the default. Any # defined in user_defined will replace the default for that # SourceConfig.name @@ -320,13 +316,13 @@ class PayoutConfig(BaseModel): """Store configuration related to payouts, payout transformation, and user payout formatting.""" - payout_format: Optional[PayoutFormatType] = Field( + payout_format: PayoutFormatType | None = Field( default=None, description=payout_format_description, examples=payout_format_examples, ) - payout_transformation: Optional[PayoutTransformation] = Field( + payout_transformation: PayoutTransformation | None = Field( default=None, description="How the BP's payout is converted to the User's payout", ) @@ -395,7 +391,7 @@ class UserCreateConfig(BaseModel): description="The smallest allowed value for the hourly user create limit.", ) - max_hourly_create_limit: Optional[NonNegativeInt] = Field( + max_hourly_create_limit: NonNegativeInt | None = Field( default=None, description="The largest allowed value for the hourly user create " "limit. If None, the hourly create limit is unconstrained.", @@ -420,11 +416,11 @@ class UserWalletConfig(BaseModel): # This field could go in supported_payout_types ---v amt: bool = Field(default=False, description="Uses Amazon Mechanical Turk") - supported_payout_types: Set["PayoutType"] = Field( + supported_payout_types: set["PayoutType"] = Field( default={PayoutType.CASH_IN_MAIL, PayoutType.TANGO, PayoutType.PAYPAL} ) - min_cashout: Optional[Decimal] = Field( + min_cashout: Decimal | None = Field( default=None, gt=0, description="Minimum cashout amount. If enabled is True and no min_cashout is " @@ -434,8 +430,8 @@ class UserWalletConfig(BaseModel): @field_serializer("supported_payout_types", when_used="json") def serialize_supported_payout_types_in_order( - self, supported_payout_types: Set["PayoutType"] - ) -> Set["PayoutType"]: + self, supported_payout_types: set["PayoutType"] + ) -> set["PayoutType"]: return set(sorted(supported_payout_types)) @field_validator("min_cashout", mode="after") @@ -470,14 +466,14 @@ class PayoutTransformationPercentArgs(BaseModel): examples=[0.5], ) - min_payout: Optional[Decimal] = Field( + min_payout: Decimal | None = Field( default=None, description="The minimum amount paid for a complete. Note: This does not " "check that the actual payout was at least this amount.", examples=[Decimal("0.50")], ) - max_payout: Optional[Decimal] = Field( + max_payout: Decimal | None = Field( default=None, description="The maximum amount paid for a complete", examples=[Decimal("5.00")], @@ -524,7 +520,7 @@ class PayoutTransformation(BaseModel): description="The name of the transformation function to use." ) - kwargs: Optional[PayoutTransformationPercentArgs] = Field( + kwargs: PayoutTransformationPercentArgs | None = Field( description="The kwargs to pass to the transformation function.", examples=[{"pct": 0.50, "max_payout": "5.00"}], default=None, @@ -549,8 +545,8 @@ class PayoutTransformation(BaseModel): self, payout: Decimal, pct: Decimal = 1, - min_payout: Decimal = 0, - max_payout: Optional[Decimal] = None, + min_payout: Decimal | None = 0, + max_payout: Decimal | None = None, ) -> Decimal: """Payout transformation for user displayed values""" if min_payout is None: @@ -567,7 +563,7 @@ class PayoutTransformation(BaseModel): return payout def payout_transformation_amt( - self, payout: Decimal, user_wallet_balance: Optional[Decimal] = None + self, payout: Decimal, user_wallet_balance: Decimal | None = None ) -> Decimal: """Payout transformation for user displayed values""" # If user_wallet_balance isn't passed, we are re-calculating this @@ -600,7 +596,7 @@ class SourceConfig(BaseModel): name: Source = Field() active: bool = Field(default=True) - banned_countries: List[CountryISOLike] = Field(default_factory=list) + banned_countries: list[CountryISOLike] = Field(default_factory=list) allow_mobile_ip: bool = Field(default=True) allow_pii_only_buyers: bool = Field( @@ -651,7 +647,7 @@ class SupplyConfig(BaseModel): model_config = ConfigDict(frozen=False, validate_assignment=True) - policies: List[SupplyPolicy] = Field(default_factory=list) + policies: list[SupplyPolicy] = Field(default_factory=list) @property def configs(self): @@ -702,13 +698,13 @@ class SupplyConfig(BaseModel): return [c for c in self.policies if c.scope == Scope.PRODUCT] @property - def global_scoped_policies_dict(self) -> Dict[Source, SupplyPolicy]: + def global_scoped_policies_dict(self) -> dict[Source, SupplyPolicy]: return {c.name: c for c in self.policies if c.scope == Scope.GLOBAL} @property def team_scoped_policies_dict( self, - ) -> Dict[str, Dict[Source, SupplyPolicy]]: + ) -> dict[str, dict[Source, SupplyPolicy]]: # str in top-level dict is the team_id d = defaultdict(dict) for c in self.team_scoped_policies: @@ -719,7 +715,7 @@ class SupplyConfig(BaseModel): @property def product_scoped_policies_dict( self, - ) -> Dict[str, Dict[Source, SupplyPolicy]]: + ) -> dict[str, dict[Source, SupplyPolicy]]: # str in top-level dict is the product_id d = defaultdict(dict) for c in self.product_scoped_policies: @@ -729,7 +725,7 @@ class SupplyConfig(BaseModel): def get_policies_for( self, product_id: str, team_id: str - ) -> Dict[Source, SupplyPolicy]: + ) -> dict[Source, SupplyPolicy]: """ Is there a config scoped to this product? If not, Is there a config scoped to this team? If not, @@ -768,20 +764,20 @@ class SupplyPolicy(SourceConfig): internally in grpc logic. """ - address: List[str] = Field(description="address for the grpc GetOpps call") + address: list[str] = Field(description="address for the grpc GetOpps call") allow_vpn: bool = Field(default=False) distribute_harmonizer_active: bool = Field(default=True) - supplier_id: Optional[str] = Field( + supplier_id: str | None = Field( default=None, description="For some inventory Sources, we may partition traffic using " "different supplier accounts instead", ) - team_ids: Optional[List[UUIDStr]] = Field(default=None) - product_ids: Optional[List[UUIDStr]] = Field(default=None) + team_ids: list[UUIDStr] | None = Field(default=None) + product_ids: list[UUIDStr] | None = Field(default=None) integration_mode: IntegrationMode = Field(default=IntegrationMode.PLATFORM) @@ -847,7 +843,7 @@ class Product(BaseModel, validate_assignment=True): examples=["1108d053e4fa47c5b0dbdcd03a7981e7"], ) - id_int: SkipJsonSchema[Optional[PositiveInt]] = Field(default=None) + id_int: SkipJsonSchema[PositiveInt | None] = Field(default=None) name: str = Field( min_length=3, @@ -868,14 +864,14 @@ class Product(BaseModel, validate_assignment=True): "be made to the Product.", ) - created: Optional[AwareDatetimeISO] = Field( + created: AwareDatetimeISO | None = Field( # TODO: make this non-nullable default=None, description="When the Product was created, this does necessarily mean " "it started to retrieve traffic at that time.", ) - team_id: Optional[UUIDStr] = Field( + team_id: UUIDStr | None = Field( # TODO: make this non-nullable default=None, examples=["b96c1209cf4a4baaa27d38082421a039"], @@ -884,7 +880,7 @@ class Product(BaseModel, validate_assignment=True): "Product", ) - business_id: Optional[UUIDStr] = Field( + business_id: UUIDStr | None = Field( default=None, examples=[uuid4().hex], description="The legal business entity or individual that is " @@ -892,7 +888,7 @@ class Product(BaseModel, validate_assignment=True): "Payments for this Product's activity.", ) - tags: Set["SupplierTag"] = Field( + tags: set["SupplierTag"] = Field( default_factory=set, description="Tags which are used to annotate supplier traffic", ) @@ -945,24 +941,24 @@ class Product(BaseModel, validate_assignment=True): # Initialization is deferred until unless it's called # (see .prebuild_***()) - balance: Optional["ProductBalances"] = Field( + balance: "ProductBalances" | None = Field( default=None, description="Product Balance" ) - payouts_total_str: Optional[str] = Field(default=None) - payouts_total: Optional[USDCent] = Field(default=None) - payouts: Optional[List["BrokerageProductPayoutEvent"]] = Field( + payouts_total_str: str | None = Field(default=None) + payouts_total: USDCent | None = Field(default=None) + payouts: list["BrokerageProductPayoutEvent"] | None = Field( default=None, description="Product Payouts. These are the ACH or Wire payments that were sent to the" "Business on behalf of this specific Product", ) - pop_financial: Optional[List["POPFinancial"]] = Field(default=None) - bp_account: Optional[LedgerAccount] = Field(default=None) + pop_financial: list["POPFinancial"] | None = Field(default=None) + bp_account: LedgerAccount | None = Field(default=None) # --- Validators --- @field_validator("harmonizer_domain", mode="before") - def harmonizer_domain_https(cls, s: Optional[str]): + def harmonizer_domain_https(cls, s: str | None): # in the db, this has no scheme. accept both with a default of https:// if s is not None and not (s.startswith("https://") or s.startswith("http://")): s = f"https://{s}" @@ -1022,7 +1018,7 @@ class Product(BaseModel, validate_assignment=True): return self.sources_config.sources @property - def sources_dict(self) -> Dict[Source, SourceConfig]: + def sources_dict(self) -> dict[Source, SourceConfig]: # This stores the same info as sources but with the keys as a Source return {x.name: x for x in self.sources} @@ -1054,12 +1050,10 @@ class Product(BaseModel, validate_assignment=True): return f"product-{self.uuid}" # --- Prefetch --- - def prefetch_bp_account(self, thl_lm: "ThlLedgerManager"): + def prefetch_bp_account(self, thl_lm: "ThlLedgerManager") -> None: account = thl_lm.get_account_or_create_bp_wallet(product=self) self.bp_account = account - return None - # --- Prebuild --- def prebuild_balance( @@ -1067,7 +1061,7 @@ class Product(BaseModel, validate_assignment=True): thl_lm: "ThlLedgerManager", ds: "GRLDatasets", client: Client, - pop_ledger: Optional["PopLedgerMerge"] = None, + pop_ledger: "PopLedgerMerge" | None = None, ) -> None: """ This returns the Product's Balances that are calculated across @@ -1129,7 +1123,6 @@ class Product(BaseModel, validate_assignment=True): "If the df is empty, we can also assume that there should be no " "transactions in the ledger." ) - return None df = df.set_index("time_idx") from generalresearch.models.thl.finance import ProductBalances @@ -1143,14 +1136,13 @@ class Product(BaseModel, validate_assignment=True): assert bal == balance.balance, "Sql and Parquet Balance inconsistent" self.balance = balance - return None def prebuild_pop_financial( self, thl_lm: "ThlLedgerManager", ds: "GRLDatasets", client: Client, - pop_ledger: Optional["PopLedgerMerge"] = None, + pop_ledger: "PopLedgerMerge" | None = None, ) -> None: """This is very similar to the Product POP Financial endpoint; however, it returns more than one item for a single time interval. This is @@ -1186,13 +1178,13 @@ class Product(BaseModel, validate_assignment=True): ) if ddf is None: self.pop_financial = [] - return None + return df = client.compute(collections=ddf, sync=True) if df.empty: self.pop_financial = [] - return None + return df = df.groupby( [pd.Grouper(key="time_idx", freq=rr.interval), "account_id"] @@ -1204,7 +1196,7 @@ class Product(BaseModel, validate_assignment=True): input_data=df, accounts=[self.bp_account] ) - return None + return def prebuild_payouts( self, @@ -1228,8 +1220,6 @@ class Product(BaseModel, validate_assignment=True): 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(self): # account = LM.get_account(qualified_name=f"{LM.currency.value}:bp_wallet:{product.id}") # @@ -1326,7 +1316,7 @@ class Product(BaseModel, validate_assignment=True): client: Client, bp_pem: "BrokerageProductPayoutEventManager", redis_config: RedisConfig, - pop_ledger: Optional[PopLedgerMerge] = None, + pop_ledger: PopLedgerMerge | None = None, ) -> None: LOG.debug(f"Product.set_cache({self.uuid=})") @@ -1352,8 +1342,6 @@ class Product(BaseModel, validate_assignment=True): rc = redis_config.create_redis_client() rc.set(name=self.cache_key, value=self.model_dump_json(), ex=ex_secs) - return None - def determine_bp_payment(self, thl_net: Decimal) -> Decimal: """ How much should we pay the BP? @@ -1369,7 +1357,6 @@ class Product(BaseModel, validate_assignment=True): return (thl_net * self.commission_pct).quantize(Decimal("0.01")) def get_payout_transformation_func(self) -> Callable: - """ """ if self.payout_config.payout_transformation is None: return lambda x: x else: @@ -1378,8 +1365,8 @@ class Product(BaseModel, validate_assignment=True): ) def calculate_user_payment( - self, bp_payout: Decimal, user_wallet_balance: Optional[Decimal] = None - ) -> Optional[Decimal]: + self, bp_payout: Decimal, user_wallet_balance: Decimal | None = None + ) -> Decimal | None: """ :param bp_payout: This is the amount we paid to the brokerage product :return: The amount that should be paid to the user @@ -1403,7 +1390,7 @@ class Product(BaseModel, validate_assignment=True): url = urlunsplit(url_split) return url - def format_payout_format(self, payout: Decimal) -> Optional[str]: + def format_payout_format(self, payout: Decimal) -> str | None: assert isinstance(payout, Decimal), "payout should be a Decimal" if self.payout_config.payout_format is None: return None @@ -1412,7 +1399,7 @@ class Product(BaseModel, validate_assignment=True): # --- ORM --- - def model_dump_mysql(self, *args, **kwargs) -> Dict[str, Any]: + def model_dump_mysql(self, *args, **kwargs) -> dict[str, Any]: d = self.model_dump(mode="json", *args, **kwargs) if "created" in d: diff --git a/generalresearch/models/thl/profiling/marketplace.py b/generalresearch/models/thl/profiling/marketplace.py index 2dd7028..027aa4c 100644 --- a/generalresearch/models/thl/profiling/marketplace.py +++ b/generalresearch/models/thl/profiling/marketplace.py @@ -1,7 +1,9 @@ +from __future__ import annotations + from abc import ABC, abstractmethod from datetime import datetime, timezone from functools import cached_property -from typing import Any, Dict, Optional, Set, Tuple +from typing import Any from pydantic import BaseModel, ConfigDict, Field, PositiveInt, computed_field @@ -37,7 +39,7 @@ class MarketplaceQuestion(BaseModel, ABC): # Refers to a Category that we annotate. The info is stored in different # dbs, so it may not be possible to retrieve the Category from the id, # so we just store the id here. - category_id: Optional[UUIDStr] = Field(default=None) + category_id: UUIDStr | None = Field(default=None) # # This doesn't work # @property @@ -58,7 +60,7 @@ class MarketplaceQuestion(BaseModel, ABC): return f"{self.source.value}:{self.internal_id}" @property - def _key(self) -> Tuple[str, CountryISOLike, LanguageISOLike]: + def _key(self) -> tuple[str, CountryISOLike, LanguageISOLike]: """This uniquely identifies a question in a locale. There is a unique index on this in the db. e.g. (question_id, country_iso, language_iso) """ @@ -68,7 +70,7 @@ class MarketplaceQuestion(BaseModel, ABC): def to_upk_question(self): ... @computed_field - def num_options(self) -> Optional[int]: + def num_options(self) -> int | None: return len(self.options) if self.options is not None else None def __hash__(self): @@ -96,7 +98,7 @@ class MarketplaceUserQuestionAnswer(BaseModel): # "anonymous" users, which are represented by a list of question answers # not associated with an actual user. No default b/c we must explicitly # set the field to None. - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32) + user_id: PositiveInt | None = Field(lt=MAX_INT32) question_id: str = Field() @@ -104,7 +106,7 @@ class MarketplaceUserQuestionAnswer(BaseModel): # these are fetched from the db for use in yield-management, we read this # field from the marketplace's question table. # This should be overloaded in each implementation !!! - question_type: Optional[str] = Field(default=None) + question_type: str | None = Field(default=None) # This may be a pipe-separated string if the question_type is multi. Regex # means any chars except capital letters @@ -116,14 +118,14 @@ class MarketplaceUserQuestionAnswer(BaseModel): language_iso: LanguageISO = Field(frozen=True) @cached_property - def options_ids(self) -> Set[str]: + def options_ids(self) -> set[str]: return set(self.option_id.split("|")) @property def pre_code(self) -> str: return self.option_id - def to_mysql(self) -> Dict[str, Any]: + def to_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"question_type"}) d["created"] = self.created.replace(tzinfo=None) return d diff --git a/generalresearch/models/thl/profiling/question.py b/generalresearch/models/thl/profiling/question.py index 72115fe..3e2984a 100644 --- a/generalresearch/models/thl/profiling/question.py +++ b/generalresearch/models/thl/profiling/question.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Dict, Optional, Tuple +from typing import Any from pydantic import ( BaseModel, @@ -21,20 +21,20 @@ from generalresearch.models.thl.profiling.upk_question import UpkQuestion class Question(BaseModel): model_config = ConfigDict(validate_assignment=True) - id: Optional[UUIDStr] = Field(default=None, alias="question_id") + id: UUIDStr | None = Field(default=None, alias="question_id") # ISO 3166-1 alpha-2 (two-letter codes, lowercase) country_iso: CountryISOLike = Field() # 3-char ISO 639-2/B, lowercase language_iso: LanguageISOLike = Field() - property_code: Optional[str] = Field( + property_code: str | None = Field( default=None, description="What marketplace question this question links to", pattern=r"^[a-z]{1,2}\:.*", ) data: UpkQuestion = Field() is_live: bool = Field() - custom: Dict[str, Any] = Field(default_factory=dict) + custom: dict[str, Any] = Field(default_factory=dict) last_updated: AwareDatetimeISO = Field() @computed_field @@ -42,5 +42,5 @@ class Question(BaseModel): def md5sum(self) -> str: return self.data.md5sum - def validate_question_answer(self, answer: Tuple[str, ...]) -> Tuple[bool, str]: + def validate_question_answer(self, answer: tuple[str, ...]) -> tuple[bool, str]: return self.data.validate_question_answer(answer=answer) diff --git a/generalresearch/models/thl/profiling/upk_property.py b/generalresearch/models/thl/profiling/upk_property.py index 9eede95..9e78692 100644 --- a/generalresearch/models/thl/profiling/upk_property.py +++ b/generalresearch/models/thl/profiling/upk_property.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from enum import Enum from functools import cached_property -from typing import Dict, List, Optional from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter @@ -34,7 +35,7 @@ class Cardinality(str, Enum, metaclass=ReprEnumMeta): class UpkItem(BaseModel): id: UUIDStr = Field(examples=["497b1fedec464151b063cd5367643ffa"]) label: str = Field(max_length=255, examples=["high_school_completion"]) - description: Optional[str] = Field( + description: str | None = Field( max_length=1024, examples=["Completed high school"], default=None ) @@ -78,17 +79,17 @@ class UpkProperty(BaseModel): "the Nudge API.", ) - allowed_items: Optional[List[UpkItem]] = Field(default=None) + allowed_items: list[UpkItem] | None = Field(default=None) - categories: List[Category] = Field(default_factory=list) + categories: list[Category] = Field(default_factory=list) @cached_property - def allowed_items_by_label(self) -> Dict[str, UpkItem]: + def allowed_items_by_label(self) -> dict[str, UpkItem]: return {i.label: i for i in self.allowed_items} @cached_property - def allowed_items_by_id(self) -> Dict[UUIDStr, UpkItem]: + def allowed_items_by_id(self) -> dict[UUIDStr, UpkItem]: return {i.id: i for i in self.allowed_items} -ProfilingInfo = TypeAdapter(List[UpkProperty]) +ProfilingInfo = TypeAdapter(list[UpkProperty]) diff --git a/generalresearch/models/thl/profiling/upk_question.py b/generalresearch/models/thl/profiling/upk_question.py index 5c908b0..307bc33 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, Dict, List, Literal, Optional, Set, Tuple, Union +from typing import Any, List, Literal, Union from pydantic import ( BaseModel, @@ -24,21 +24,21 @@ from generalresearch.models.thl.category import Category class UPKImportance(BaseModel): - task_count: Optional[int] = Field( + task_count: int | None = Field( ge=0, default=None, examples=[47], description="The number of live Tasks that use this UPK Question", ) - task_score: Optional[float] = Field( + task_score: float | None = Field( ge=0, default=None, examples=[0.11175522477414712], description="GRL's internal ranked score for the UPK Question", ) - marketplace_task_count: Optional[Dict[Source, NonNegativeInt]] = Field( + marketplace_task_count: dict[Source, NonNegativeInt] | None = Field( default=None, examples=[{Source.DYNATA: 23, Source.SPECTRUM: 24}], description="The number of live Tasks that use this UPK Question per marketplace", @@ -80,14 +80,14 @@ class UpkQuestionChoice(BaseModel): order: NonNegativeInt = Field() # Allows you to group answer choices together (used for display or extra logic) - group: Optional[int] = Field(default=None) + group: int | None = Field(default=None) exclusive: bool = Field( default=False, description="If answer is exclusive, it can be the only option selected", ) - importance: Optional[UPKImportance] = Field(default=None) + importance: UPKImportance | None = Field(default=None) def __hash__(self): # We don't know the question ID!! Unique within a question only! @@ -151,7 +151,7 @@ class UpkQuestionConfigurationMC(BaseModel): exclude=True, default=UpkQuestionType.MULTIPLE_CHOICE ) - max_select: Optional[int] = Field(gt=0, default=None) + max_select: int | None = Field(gt=0, default=None) class UpkQuestionConfigurationTE(BaseModel): @@ -163,7 +163,7 @@ class UpkQuestionConfigurationTE(BaseModel): ) # Sets input form attribute; not the same as regex validation - max_length: Optional[PositiveInt] = Field( + max_length: PositiveInt | None = Field( default=None, description="Maximum str length of any input. Meant as an easy, non" "regex based check.", @@ -171,7 +171,7 @@ class UpkQuestionConfigurationTE(BaseModel): # The text input box must contain this number of chars before submission # is allowed - min_length: Optional[PositiveInt] = Field( + min_length: PositiveInt | None = Field( default=None, description="Minimum str length of any input. Meant as an easy, non" "regex based check.", @@ -195,25 +195,25 @@ class UpkQuestionConfigurationSLIDER(BaseModel): ) # TODO: constraints. we don't have any of these so not wasting time on this - slider_min: Optional[float] = Field(default=None) - slider_max: Optional[float] = Field(default=None) - slider_start: Optional[float] = Field(default=None) - slider_step: Optional[float] = Field(default=None) + slider_min: float | None = Field(default=None) + slider_max: float | None = Field(default=None) + slider_start: float | None = Field(default=None) + slider_step: float | None = Field(default=None) class UpkQuestionValidation(BaseModel): model_config = ConfigDict(frozen=True) # --- UpkQuestionType.TEXT_ENTRY Options --- - patterns: Optional[List[PatternValidation]] = Field(min_length=1) + patterns: list[PatternValidation] | None = Field(min_length=1) -SelectorType = Union[ - UpkQuestionSelectorMC, - UpkQuestionSelectorTE, - UpkQuestionSelectorSLIDER, - UpkQuestionSelectorHIDDEN, -] +SelectorType = ( + UpkQuestionSelectorMC + | UpkQuestionSelectorTE + | UpkQuestionSelectorSLIDER + | UpkQuestionSelectorHIDDEN +) Configuration = Annotated[ Union[ UpkQuestionConfigurationMC, @@ -276,11 +276,11 @@ class UpkQuestion(BaseModel): ) # The id is globally unique - id: Optional[UUIDStr] = Field(default=None, alias="question_id") + id: UUIDStr | None = Field(default=None, alias="question_id") # The format is "{Source}:{question_id}" where Source is 1 or 2 chars, and # question_id is the marketplace's ID for this question. - ext_question_id: Optional[str] = Field( + ext_question_id: str | None = Field( default=None, description="what marketplace question this question links to", pattern=r"^[a-z]{1,2}\:.*", @@ -301,25 +301,25 @@ class UpkQuestion(BaseModel): # Don't set a min_length=1 here. We'll allow this to be created, but it # won't be askable with empty choices. - choices: Optional[List[UpkQuestionChoice]] = Field(default=None) + choices: list[UpkQuestionChoice] | None = Field(default=None) selector: SelectorType = Field() - configuration: Optional[Configuration] = Field(default=None) - validation: Optional[UpkQuestionValidation] = Field(default=None) - importance: Optional[UPKImportance] = Field(default=None) + configuration: Configuration | None = Field(default=None) + validation: UpkQuestionValidation | None = Field(default=None) + importance: UPKImportance | None = Field(default=None) - categories: List[Category] = Field( + categories: list[Category] = Field( default_factory=list, description="Categories associated with this question", ) - explanation_template: Optional[str] = Field( + explanation_template: str | None = Field( description="Human-readable template for explaining how a user's answer to this question affects eligibility", examples=[ "The company that administers your employer's retirement plan is {answer}." ], default=None, ) - explanation_fragment_template: Optional[str] = Field( + explanation_fragment_template: str | None = Field( description="A very short, natural-language explanation fragment that can be combined with others into a single sentence", examples=["whose employer's retirement plan is {answer}"], default=None, @@ -336,7 +336,7 @@ class UpkQuestion(BaseModel): return self.country_iso + "_" + self.language_iso @property - def source(self) -> Optional[Source]: + def source(self) -> Source | None: if self.ext_question_id: return Source(self.ext_question_id.split(":", 1)[0]) @@ -348,7 +348,7 @@ class UpkQuestion(BaseModel): @model_validator(mode="before") @classmethod - def check_configuration_type(cls, data: Dict[str, Any]) -> Dict[str, Any]: + def check_configuration_type(cls, data: dict[str, Any]) -> dict[str, Any]: # The model knows what the type of Configuration to grab depending on # the key 'type' which it expects inside the configuration object. # Here, we grab the type from the top-level model instead. @@ -441,8 +441,8 @@ class UpkQuestion(BaseModel): @field_validator("choices") @classmethod def validate_choices( - cls, choices: Optional[List[UpkQuestionChoice]] - ) -> Optional[List[UpkQuestionChoice]]: + cls, choices: list[UpkQuestionChoice] | None + ) -> list[UpkQuestionChoice] | None: if choices: ids = {x.id for x in choices} assert len(ids) == len(choices), "choices.id must be unique" @@ -507,7 +507,7 @@ class UpkQuestion(BaseModel): d.update(d.pop("importance", {})) return d - def validate_question_answer(self, answer: Tuple[str, ...]) -> Tuple[bool, str]: + def validate_question_answer(self, answer: tuple[str, ...]) -> tuple[bool, str]: """ Returns (is_valid, error_message). """ @@ -518,7 +518,7 @@ class UpkQuestion(BaseModel): else: return True, "" - def _validate_question_answer(self, answer: Tuple[str, ...]) -> None: + def _validate_question_answer(self, answer: tuple[str, ...]) -> None: """ If the question is MC, validate: - validate selector SA vs MA (1 selected vs >1 selected) @@ -578,24 +578,24 @@ class UpkQuestion(BaseModel): class UpkQuestionOut(UpkQuestion): - choices: Optional[List[UpkQuestionChoiceOut]] = Field(default=None) + choices: list[UpkQuestionChoiceOut] | None = Field(default=None) # Return both importance top-level model and extracted keys for now. # Eventually deprecate one way. - task_count: Optional[int] = Field( + task_count: int | None = Field( ge=0, default=None, examples=[47], description="The number of live Tasks that use this UPK Question", ) - task_score: Optional[float] = Field( + task_score: float | None = Field( ge=0, default=None, examples=[0.11175522477414712], description="GRL's internal ranked score for the UPK Question", ) - marketplace_task_count: Optional[Dict[Source, NonNegativeInt]] = Field( + marketplace_task_count: dict[Source, NonNegativeInt] | None = Field( default=None, examples=[{Source.DYNATA: 23, Source.SPECTRUM: 24}], description="The number of live Tasks that use this UPK Question per marketplace", @@ -649,7 +649,7 @@ def trim_options(q: UpkQuestion, max_options: int = 50) -> UpkQuestion: if q.ext_question_id.startswith("gr:") or q.ext_question_id.startswith("g:"): return q - special_choices: Set[UpkQuestionChoice] = { + special_choices: set[UpkQuestionChoice] = { c for c in q.choices if option_is_catch_all(c) } diff --git a/generalresearch/models/thl/profiling/upk_question_answer.py b/generalresearch/models/thl/profiling/upk_question_answer.py index 2eb52e1..0024e68 100644 --- a/generalresearch/models/thl/profiling/upk_question_answer.py +++ b/generalresearch/models/thl/profiling/upk_question_answer.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from datetime import datetime, timezone -from typing import Any, Dict, Optional, Union +from typing import Any from uuid import uuid4 from pydantic import ( @@ -30,12 +32,12 @@ class UpkQuestionAnswer(BaseModel): user_id: PositiveInt = Field(lt=MAX_INT32) - question_id: Optional[UUIDStr] = Field( + question_id: UUIDStr | None = Field( examples=[uuid4().hex], description="The ID of the question that was asked in order to determine this", default=None, ) - session_id: Optional[UUIDStr] = Field( + session_id: UUIDStr | None = Field( examples=[uuid4().hex], description="The thl_session in which the question was asked", default=None, @@ -64,23 +66,23 @@ class UpkQuestionAnswer(BaseModel): # 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). - item_id: Optional[UUIDStr] = Field( + item_id: UUIDStr | None = Field( default=None, examples=["497b1fedec464151b063cd5367643ffa"] ) - item_label: Optional[str] = Field( + item_label: str | None = Field( default=None, max_length=255, examples=["high_school_completion"] ) - value_text: Optional[str] = Field( + value_text: str | None = Field( default=None, max_length=1024, ) - value_num: Optional[float] = Field( + value_num: float | None = Field( default=None, ) @computed_field @property - def value(self) -> Optional[Union[str, float]]: + def value(self) -> str | float | None: if self.prop_type == PropertyType.UPK_ITEM: return self.item_label elif self.prop_type == PropertyType.UPK_TEXT: @@ -109,7 +111,7 @@ class UpkQuestionAnswer(BaseModel): return self - def model_dump_mysql(self) -> Dict[str, Any]: + def model_dump_mysql(self) -> dict[str, Any]: d = self.model_dump(mode="json") d["created"] = self.created return d diff --git a/generalresearch/models/thl/profiling/user_info.py b/generalresearch/models/thl/profiling/user_info.py index 15197c9..32af704 100644 --- a/generalresearch/models/thl/profiling/user_info.py +++ b/generalresearch/models/thl/profiling/user_info.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from __future__ import annotations from pydantic import BaseModel, ConfigDict, Field from pydantic.json_schema import SkipJsonSchema @@ -14,11 +14,11 @@ from generalresearch.models.thl.user import User class UserProfileKnowledgeAnswer(BaseModel): # Returns {id, label, translation} when the prop_type is an item, # and only {value} if it's a string/text (such as for postalcode) - id: Optional[str] = Field(default=None) - label: Optional[str] = Field(default=None) - translation: Optional[str] = Field(default=None) + id: str | None = Field(default=None) + label: str | None = Field(default=None) + translation: str | None = Field(default=None) - value: Optional[str] = Field(default=None) + value: str | None = Field(default=None) class UserProfileKnowledge(BaseModel): @@ -26,7 +26,7 @@ class UserProfileKnowledge(BaseModel): property_label: str = Field() translation: str = Field() - answer: List[UserProfileKnowledgeAnswer] = Field(default_factory=list) + answer: list[UserProfileKnowledgeAnswer] = Field(default_factory=list) created: AwareDatetimeISO = Field( description="When the User submitted this Profiling data" @@ -44,7 +44,7 @@ class MarketProfileKnowledge(BaseModel): question_id: str = Field(examples=["gender", "1843", "gender_plus"]) - answer: List[str] = Field( + answer: list[str] = Field( default_factory=list, examples=[["male"], ["7657644"], ["1"]] ) @@ -67,10 +67,10 @@ class MarketProfileKnowledge(BaseModel): class UserInfo(BaseModel): model_config = ConfigDict() - user: SkipJsonSchema[Optional[User]] = Field(exclude=True, default=None) + user: SkipJsonSchema[User | None] = Field(exclude=True, default=None) - user_profile_knowledge: List[UserProfileKnowledge] = Field(default_factory=list) + user_profile_knowledge: list[UserProfileKnowledge] = Field(default_factory=list) - marketplace_profile_knowledge: List[MarketProfileKnowledge] = Field( + marketplace_profile_knowledge: list[MarketProfileKnowledge] = Field( default_factory=list ) diff --git a/generalresearch/models/thl/profiling/user_question_answer.py b/generalresearch/models/thl/profiling/user_question_answer.py index b325583..0d3e37c 100644 --- a/generalresearch/models/thl/profiling/user_question_answer.py +++ b/generalresearch/models/thl/profiling/user_question_answer.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import json from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterator, Literal, Optional, Tuple, Union +from typing import Any, Iterator, Literal from pydantic import ( BaseModel, @@ -23,30 +25,30 @@ class UserQuestionAnswer(BaseModel): model_config = ConfigDict(validate_assignment=True) - user_id: Optional[PositiveInt] = Field(lt=MAX_INT32, default=None) + user_id: PositiveInt | None = Field(lt=MAX_INT32, default=None) question_id: UUIDStr = Field() - answer: Tuple[str, ...] = Field() + answer: tuple[str, ...] = Field() timestamp: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc) ) - country_iso: Union[CountryISO, Literal["xx"]] = Field() - language_iso: Union[LanguageISO, Literal["xxx"]] = Field() + country_iso: CountryISO | Literal["xx"] = Field() + language_iso: LanguageISO | Literal["xxx"] = Field() # Store a property code associated with this question_id. e.g. "gr:hispanic" or "d:192" property_code: str = Field() # Stores any question answers that are calculated from this answer - calc_answers: Optional[Dict[str, Tuple[str, ...]]] = Field(default=None) + calc_answers: dict[str, tuple[str, ...]] | None = Field(default=None) @field_validator("calc_answers") - def sorted_calc_answers(cls, calc_answers) -> Optional[Dict[str, Tuple[str, ...]]]: + def sorted_calc_answers(cls, calc_answers) -> dict[str, tuple[str, ...]] | None: if calc_answers is None: return None return {k: tuple(sorted(v)) for k, v in calc_answers.items()} @field_validator("calc_answers") - def validate_keys(cls, calc_answers) -> Optional[Dict[str, Tuple[str, ...]]]: + def validate_keys(cls, calc_answers) -> dict[str, tuple[str, ...]] | None: if calc_answers is None: return None @@ -55,7 +57,7 @@ class UserQuestionAnswer(BaseModel): ), "calc_answers expects the keys to be in format source:question_code" return calc_answers - def model_dump_mysql(self, session_id: Optional[str] = None) -> Dict[str, Any]: + def model_dump_mysql(self, session_id: str | None = None) -> dict[str, Any]: d = self.model_dump(mode="json", exclude={"calc_answers", "timestamp"}) d["answer"] = json.dumps(self.answer) # Note naming inconsistency here: calc_answer/s @@ -84,7 +86,7 @@ class UserQuestionAnswer(BaseModel): def __hash__(self) -> int: return hash((self.question_id, self.answer, self.timestamp)) - def validate_question_answer(self, question: UpkQuestion) -> Tuple[bool, str]: + def validate_question_answer(self, question: UpkQuestion) -> tuple[bool, str]: """ Returns (is_valid, error_message). """ @@ -143,7 +145,7 @@ class MarketplaceResearchProfileQuestion(BaseModel): description="# the question id/code on the marketplace", min_length=1 ) source: Source = Field() # the one or two-letter marketplace code - answer: Tuple[str, ...] = Field(min_length=1) + answer: tuple[str, ...] = Field(min_length=1) timestamp: AwareDatetimeISO = Field() country_iso: CountryISO = Field() language_iso: LanguageISO = Field() diff --git a/generalresearch/models/thl/report_task.py b/generalresearch/models/thl/report_task.py index 1abd330..d29599d 100644 --- a/generalresearch/models/thl/report_task.py +++ b/generalresearch/models/thl/report_task.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import random from collections import defaultdict -from typing import Collection, List, Optional +from collections.abc import Collection from pydantic import BaseModel, ConfigDict, Field @@ -21,7 +23,7 @@ REPORT_PRIORITY = defaultdict( def prioritize_report_values( report_values: Collection[ReportValue], -) -> Optional[ReportValue]: +) -> ReportValue | None: if not report_values: return None report_values = list(set(report_values)) @@ -39,7 +41,7 @@ class ReportTask(BaseModel): examples=["app-user-9329ebd"], ) - reasons: List[ReportValue] = Field( + reasons: list[ReportValue] = Field( description=ReportValue.as_openapi_with_value_descriptions(), examples=[[3, 4]], default_factory=list, diff --git a/generalresearch/models/thl/session.py b/generalresearch/models/thl/session.py index 74ed5eb..fbaf0e7 100644 --- a/generalresearch/models/thl/session.py +++ b/generalresearch/models/thl/session.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import json import logging from datetime import datetime, timedelta, timezone from decimal import Decimal -from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Annotated, Any from uuid import uuid4 from pydantic import ( @@ -64,7 +66,7 @@ class WallBase(BaseModel): uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) source: Source - buyer_id: Optional[str] = Field(default=None, max_length=32) + 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( @@ -74,27 +76,27 @@ class WallBase(BaseModel): # 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 # on instantiation - survey_id: Optional[str] = Field(max_length=32, default=None) - cpi: Optional[Decimal] = Field(lt=1000, ge=0, default=None) + survey_id: str | None = Field(max_length=32, default=None) + cpi: Decimal | None = Field(lt=1000, ge=0, default=None) # Gets set when a wall is "finished" - finished: Optional[AwareDatetimeISO] = Field(default=None) - status: Optional[Status] = None - status_code_1: Optional[StatusCode1] = None - status_code_2: Optional[WallStatusCode2] = None + finished: AwareDatetimeISO | None = Field(default=None) + status: Status | None = None + status_code_1: StatusCode1 | None = None + status_code_2: WallStatusCode2 | None = None - ext_status_code_1: Optional[str] = Field(default=None, max_length=32) - ext_status_code_2: Optional[str] = Field(default=None, max_length=32) - ext_status_code_3: Optional[str] = Field(default=None, max_length=32) + ext_status_code_1: str | None = Field(default=None, max_length=32) + ext_status_code_2: str | None = Field(default=None, max_length=32) + ext_status_code_3: str | None = Field(default=None, max_length=32) - report_value: Optional[ReportValue] = None - report_notes: Optional[str] = Field(default=None, max_length=255) + report_value: ReportValue | None = None + report_notes: str | None = Field(default=None, max_length=255) # This is the most recent reconciliation status of the wall event. # Possible values: 'ac' (adjusted to complete), 'af' (adj to fail) # If a wall gets adjusted and adjusted back to its original status, the # adjusted_status = None - adjusted_status: Optional[WallAdjustedStatus] = None + adjusted_status: WallAdjustedStatus | None = None # This is not really used, it is only important if the requested CPI # doesn't match the adjusted amount, which shouldn't happen as no @@ -103,11 +105,11 @@ class WallBase(BaseModel): # - If adjusted_status = 'ac': adjusted_cpi is the amount paid (should # equal the `cpi`) # - If adjusted_status = 'af': adjusted_cpi is 0.00 - adjusted_cpi: Optional[Decimal] = Field(default=None, lt=1000, ge=0) + adjusted_cpi: Decimal | None = Field(default=None, lt=1000, ge=0) # This timestamp gets updated every time there is an adjustment. Even if # we flip-flop, this will be set (and adjusted_status will be None). - adjusted_timestamp: Optional[AwareDatetimeISO] = Field(default=None) + adjusted_timestamp: AwareDatetimeISO | None = Field(default=None) # --- Validation --- @@ -261,19 +263,17 @@ class WallBase(BaseModel): self.model_config["validate_assignment"] = True self.__class__.model_validate(self) - return None - def finish( self, status: Status, status_code_1: StatusCode1, - status_code_2: Optional[WallStatusCode2] = None, - finished: Optional[datetime] = None, - ext_status_code_1: Optional[str] = None, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, - survey_id: Optional[str] = None, - cpi: Optional[Decimal] = None, + status_code_2: WallStatusCode2 | None = None, + finished: datetime | None = None, + ext_status_code_1: str | None = None, + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, + survey_id: str | None = None, + cpi: Decimal | None = None, ) -> None: # This is just used in tests at the moment. This needs to be adjusted. @@ -296,14 +296,12 @@ class WallBase(BaseModel): if cpi is not None: self.cpi = cpi - return None - def annotate_status_codes( self, ext_status_code_1: str, - ext_status_code_2: Optional[str] = None, - ext_status_code_3: Optional[str] = None, - finished: Optional[datetime] = None, + ext_status_code_2: str | None = None, + ext_status_code_3: str | None = None, + finished: datetime | None = None, ) -> None: # This should be called by the wall manager in order to actually update db from generalresearch import wall_status_codes @@ -326,8 +324,6 @@ class WallBase(BaseModel): finished=finished, ) - return None - def is_soft_fail(self) -> bool: from generalresearch import wall_status_codes @@ -371,8 +367,8 @@ class WallBase(BaseModel): def report( self, report_value: ReportValue, - report_notes: Optional[str] = None, - report_timestamp: Optional[AwareDatetime] = None, + report_notes: str | None = None, + report_timestamp: AwareDatetime | None = None, ) -> None: """When a wall event is reported: @@ -454,7 +450,7 @@ class Wall(WallBase): d = self.model_dump(mode="json", exclude={"elapsed"}) return json.dumps(d) - def model_dump_mysql(self, *args, **kwargs) -> Dict: + def model_dump_mysql(self, *args, **kwargs) -> dict: # Generate a dictionary representation of the model, with special handling for datetimes d = self.model_dump(mode="json", exclude={"elapsed"}, *args, **kwargs) d["started"] = self.started.replace(tzinfo=None) @@ -468,27 +464,27 @@ class Wall(WallBase): class WallOut(WallBase): # These get serialized to the enum name instead of the int value (for ease in UI) - status_code_1: Optional[Annotated[StatusCode1, EnumNameSerializer]] = Field( + status_code_1: Annotated[StatusCode1, EnumNameSerializer] | None = Field( default=None, examples=[StatusCode1.COMPLETE.name], description=StatusCode1.as_openapi_with_value_descriptions_name(), ) - status_code_2: Optional[Annotated[WallStatusCode2, EnumNameSerializer]] = Field( + status_code_2: Annotated[WallStatusCode2, EnumNameSerializer] | None = Field( default=None, examples=[None], description=WallStatusCode2.as_openapi_with_value_descriptions_name(), ) # Exclude these 3 fields - cpi: Optional[Decimal] = Field(lt=1000, ge=0, default=None, exclude=True) - req_cpi: Optional[Decimal] = Field( + cpi: Decimal | None = Field(lt=1000, ge=0, default=None, exclude=True) + req_cpi: Decimal | None = Field( decimal_places=5, lt=1000, ge=0, default=None, exclude=True ) - adjusted_cpi: Optional[Decimal] = Field(lt=1000, ge=0, default=None, exclude=True) + adjusted_cpi: Decimal | None = Field(lt=1000, ge=0, default=None, exclude=True) # user_cpi is serialized to integer cents!!! - user_cpi: Optional[Decimal] = Field( + user_cpi: Decimal | None = Field( lt=1000, ge=0, default=None, @@ -500,7 +496,7 @@ class WallOut(WallBase): examples=[123], ) - user_cpi_string: Optional[str] = Field( + user_cpi_string: str | None = Field( default=None, description="If a payout transformation is configured on this account, " "this is the amount to display to the user", @@ -588,7 +584,7 @@ class Session(BaseModel): # id will be None until db_create is called (or if this is instantiated # from an existing session) - id: Optional[int] = None + id: int | None = None uuid: UUIDStr = Field(default_factory=lambda: uuid4().hex) user: User started: AwareDatetimeISO = Field( @@ -599,25 +595,23 @@ class Session(BaseModel): # store the 4 fields: loi_min, loi_max, user_payout_min, user_payout_max # in the db, but there may be other metadata associated with the bucket # that is cached, such as the category. - clicked_bucket: Optional[Bucket] = Field(default=None) + clicked_bucket: Bucket | None = Field(default=None) - country_iso: Optional[str] = Field( - default=None, max_length=2, pattern=r"^[a-z]{2}$" - ) - device_type: Optional[DeviceType] = Field(default=None) - ip: Optional[IPvAnyAddressStr] = Field(default=None) + country_iso: str | None = Field(default=None, max_length=2, pattern=r"^[a-z]{2}$") + device_type: DeviceType | None = Field(default=None) + ip: IPvAnyAddressStr | None = Field(default=None) - url_metadata: Optional[Dict[str, str]] = Field(default=None) + url_metadata: dict[str, str] | None = Field(default=None) # Below here shouldn't be set upon initialization, or directly. - wall_events: List[Wall] = Field(default_factory=list) + wall_events: list[Wall] = Field(default_factory=list) # Gets set when a session is "finished" - finished: Optional[AwareDatetimeISO] = Field(default=None) + finished: AwareDatetimeISO | None = Field(default=None) - status: Optional[Status] = None - status_code_1: Optional[StatusCode1] = None - status_code_2: Optional[SessionStatusCode2] = None + status: Status | None = None + status_code_1: StatusCode1 | None = None + status_code_2: SessionStatusCode2 | None = None # There are two scenarios. Let's say the user payout transformation is # 40% and this session pays out $1. @@ -637,25 +631,25 @@ class Session(BaseModel): # going to the BP and ($0.40) to the user, and if the wallet is disabled, # then the whole $1 goes to the BP and $0 to the user, but the $0.40 value # is saved, so it can be displayed in the task status endpoint. - payout: Optional[Decimal] = Field(default=None, lt=1000, ge=0) - user_payout: Optional[Decimal] = Field(default=None, lt=1000, ge=0) + payout: Decimal | None = Field(default=None, lt=1000, ge=0) + user_payout: Decimal | None = Field(default=None, lt=1000, ge=0) # This is the most recent reconciliation status of the session. Generally, # we would adjust this if the last survey in the session was adjusted # from complete to incomplete. If any survey in the session was adjusted # from fail -> complete (and the user didn't already get a complete) # we'll adjust this to a complete. - adjusted_status: Optional[SessionAdjustedStatus] = None + adjusted_status: SessionAdjustedStatus | None = None # If adjusted_status = 'ac': payout = 0 and adjusted_payout is the amount paid # If adjusted_status = 'af': payout = the amount paid, adjusted_payout is 0.00 # (the `payout` never changed, only the adjusted_payout can change). - adjusted_payout: Optional[Decimal] = Field(default=None, lt=1000, ge=0) - adjusted_user_payout: Optional[Decimal] = Field(default=None, lt=1000, ge=0) + adjusted_payout: Decimal | None = Field(default=None, lt=1000, ge=0) + adjusted_user_payout: Decimal | None = Field(default=None, lt=1000, ge=0) # This timestamp gets updated every time there is an adjustment (even if # there are flip-flops). - adjusted_timestamp: Optional[AwareDatetimeISO] = Field(default=None) + adjusted_timestamp: AwareDatetimeISO | None = Field(default=None) # --- Validation --- @@ -757,7 +751,7 @@ class Session(BaseModel): @field_validator("wall_events") @classmethod - def check_wall_events(cls, wall_events: List[Wall]): + def check_wall_events(cls, wall_events: list[Wall]): # Note: this can't work on modifications as pydantic/python doesn't # know if a list is mutated. We have to run it manually, or hide # the self.wall_events attr and wrap all access @@ -831,7 +825,7 @@ class Session(BaseModel): def model_dump_mysql( self, *args, **kwargs - ) -> Dict[str, Union[str, int, datetime, float, None]]: + ) -> dict[str, str | int | datetime | float | None]: # Generate a dictionary representation of the model, with special # handling for datetimes, and nested models such as User & Bucket @@ -893,7 +887,7 @@ class Session(BaseModel): last_wall.status = Status.TIMEOUT self.status = Status.TIMEOUT - def determine_session_status(self) -> Tuple[Status, StatusCode1]: + def determine_session_status(self) -> tuple[Status, StatusCode1]: """Given a list of wall events, determine what the session status should be. If this is called, it is because the Session is *over*, or it has timed out. @@ -989,8 +983,8 @@ class Session(BaseModel): def determine_payments( self, - thl_ledger_manager: Optional["ThlLedgerManager"] = None, - ) -> Tuple[Decimal, Decimal, Decimal, Optional[Decimal]]: + thl_ledger_manager: "ThlLedgerManager" | None = None, + ) -> tuple[Decimal, Decimal, Decimal, Decimal | None]: # How much we should get paid by the MPs for all completes in this # session (usually 0 or 1 completes) thl_net: Decimal = Decimal( @@ -1012,7 +1006,7 @@ class Session(BaseModel): assert thl_ledger_manager is not None amt = thl_ledger_manager.get_user_wallet_balance(user=self.user) user_wallet_balance = Decimal(amt / 100).quantize(Decimal("0.01")) - user_pay: Optional[Decimal] = product.calculate_user_payment( + user_pay: Decimal | None = product.calculate_user_payment( bp_pay, user_wallet_balance=user_wallet_balance ) @@ -1034,7 +1028,7 @@ class Session(BaseModel): def determine_new_status_and_payouts( self, - ) -> Tuple[Status, Decimal, Optional[Decimal]]: + ) -> tuple[Status, Decimal, Decimal | None]: """Session is adjusted any time one of the wall events is. Assuming status adjustments happened on a session's wall events. Calculate if any status changes are need to the session. @@ -1229,7 +1223,7 @@ class Session(BaseModel): else: return self.payout or Decimal(0) - def get_user_payout_after_adjustment(self) -> Optional[Decimal]: + def get_user_payout_after_adjustment(self) -> Decimal | None: if self.adjusted_status is not None: return self.adjusted_user_payout else: @@ -1270,12 +1264,12 @@ def check_adjusted_status_consistent( def check_adjusted_status_wall_consistent( status: Status, - cpi: Optional[Decimal] = None, - adjusted_status: Optional[WallAdjustedStatus] = None, - adjusted_cpi: Optional[Decimal] = None, - new_adjusted_status: Optional[WallAdjustedStatus] = None, - new_adjusted_cpi: Optional[Decimal] = None, -) -> Tuple[bool, str]: + cpi: Decimal | None = None, + adjusted_status: WallAdjustedStatus | None = None, + adjusted_cpi: Decimal | None = None, + new_adjusted_status: WallAdjustedStatus | None = None, + new_adjusted_cpi: Decimal | None = None, +) -> tuple[bool, str]: """ Raises an AssertionError if inconsistent. @@ -1300,11 +1294,11 @@ def check_adjusted_status_wall_consistent( def _check_adjusted_status_wall_consistent( status: Status, - cpi: Optional[Decimal] = None, - adjusted_status: Optional[WallAdjustedStatus] = None, - adjusted_cpi: Optional[Decimal] = None, - new_adjusted_status: Optional[WallAdjustedStatus] = None, - new_adjusted_cpi: Optional[Decimal] = None, + cpi: Decimal | None = None, + adjusted_status: WallAdjustedStatus | None = None, + adjusted_cpi: Decimal | None = None, + new_adjusted_status: WallAdjustedStatus | None = None, + new_adjusted_cpi: Decimal | None = None, ) -> None: """ See check_adjusted_status_wall_consistent diff --git a/generalresearch/models/thl/soft_pair.py b/generalresearch/models/thl/soft_pair.py index b8bc1b4..6ff1165 100644 --- a/generalresearch/models/thl/soft_pair.py +++ b/generalresearch/models/thl/soft_pair.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from dataclasses import dataclass from enum import Enum -from typing import Optional, Set from generalresearch.models import Source from generalresearch.models.thl.survey.condition import ( @@ -33,14 +34,14 @@ class SoftPairResult: pair_type: SoftPairResultType source: Source survey_id: str - conditions: Optional[Set[MarketplaceCondition]] = None + conditions: set[MarketplaceCondition] | None = None @property def survey_sid(self) -> str: return self.source + ":" + self.survey_id @property - def grpc_string(self) -> Optional[str]: + def grpc_string(self) -> str | None: # This is what is expected by thl-grpc in a mp_pb2.MPOpportunityIDListSoftPairing response (grpc) if self.pair_type == SoftPairResultType.UNCONDITIONAL: return self.survey_id @@ -60,7 +61,7 @@ class SoftPairResultOut: pair_type: SoftPairResultType source: Source survey_id: str - question_ids: Optional[Set[str]] = None + question_ids: set[str] | None = None @property def survey_sid(self) -> str: diff --git a/generalresearch/models/thl/stats.py b/generalresearch/models/thl/stats.py index 9e9971d..40fb29a 100644 --- a/generalresearch/models/thl/stats.py +++ b/generalresearch/models/thl/stats.py @@ -1,4 +1,4 @@ -from typing import Optional +from __future__ import annotations from pydantic import BaseModel, Field, computed_field, model_validator @@ -13,7 +13,7 @@ class StatisticalSummary(BaseModel): min: int = Field() max: int = Field() - mean: Optional[int] = Field(default=None) + mean: int | None = Field(default=None) q1: int = Field() q2: int = Field() q3: int = Field() diff --git a/generalresearch/models/thl/survey/__init__.py b/generalresearch/models/thl/survey/__init__.py index 8f27930..b6ac740 100644 --- a/generalresearch/models/thl/survey/__init__.py +++ b/generalresearch/models/thl/survey/__init__.py @@ -1,7 +1,9 @@ +from __future__ import annotations + from abc import ABC, abstractmethod from decimal import Decimal from itertools import product -from typing import Dict, List, Optional, Set, Type +from typing import Type from more_itertools import flatten from pydantic import BaseModel, Field @@ -48,22 +50,22 @@ class MarketplaceTask(BaseModel, ABC): language_iso: LanguageISO = Field() # These should be overloaded with more specific type hints - buyer_id: Optional[str] = Field(min_length=1, max_length=32, default=None) + buyer_id: str | None = Field(min_length=1, max_length=32, default=None) # This is in seconds - bid_loi: Optional[int] = Field(default=None, le=90 * 60) - bid_ir: Optional[float] = Field(default=None, ge=0, le=1) + bid_loi: int | None = Field(default=None, le=90 * 60) + bid_ir: float | None = Field(default=None, ge=0, le=1) # This should be an "abstract field", but there is no way to do that, so # just listing it here. It should be overridden by the implementation source: Source = Field() # This should also - used_question_ids: Set[str] = Field(default_factory=set) + used_question_ids: set[str] = 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, MarketplaceCondition]] = Field(default=None) + conditions: dict[str, MarketplaceCondition] | None = Field(default=None) @property @abstractmethod @@ -80,7 +82,7 @@ class MarketplaceTask(BaseModel, ABC): @property @abstractmethod - def all_hashes(self) -> Set[str]: ... + def all_hashes(self) -> set[str]: ... @property @abstractmethod @@ -124,7 +126,7 @@ class MarketplaceTask(BaseModel, ABC): @abstractmethod def marketplace_genders( self, - ) -> Dict[Gender, Optional[MarketplaceCondition]]: + ) -> dict[Gender, MarketplaceCondition | None]: """ Mapping of generic Gender to the marketplace condition for that gender """ @@ -133,7 +135,7 @@ class MarketplaceTask(BaseModel, ABC): @property def marketplace_age_groups( self, - ) -> Dict[AgeGroup, Optional[MarketplaceCondition]]: + ) -> dict[AgeGroup, MarketplaceCondition | None]: """ Mapping of generic age groups to the marketplace condition for those ages """ @@ -147,7 +149,7 @@ class MarketplaceTask(BaseModel, ABC): } @property - def targeted_ages(self) -> Set[str]: + def targeted_ages(self) -> set[str]: assert self.conditions is not None, "conditions must be populated" cs = [self.conditions[k] for k in self.all_hashes if k in self.conditions] age_cs = [c for c in cs if c.question_id == self.age_question] @@ -166,7 +168,7 @@ class MarketplaceTask(BaseModel, ABC): return age_values @property - def targeted_age_groups(self) -> Set[AgeGroup]: + def targeted_age_groups(self) -> set[AgeGroup]: age_values = self.targeted_ages age_conditions = self.marketplace_age_groups age_targeting = set() @@ -186,7 +188,7 @@ class MarketplaceTask(BaseModel, ABC): return age_targeting @property - def targeted_genders(self) -> Set[Gender]: + def targeted_genders(self) -> set[Gender]: mp_genders = self.marketplace_genders gender_targeting = set() if mp_genders[Gender.MALE].criterion_hash in self.all_hashes: @@ -198,7 +200,7 @@ class MarketplaceTask(BaseModel, ABC): return gender_targeting @property - def demographic_targets(self) -> List[DemographicTarget]: + def demographic_targets(self) -> list[DemographicTarget]: targets = [DemographicTarget(country="*", gender="*", age_group="*")] gt = self.targeted_genders diff --git a/generalresearch/models/thl/survey/buyer.py b/generalresearch/models/thl/survey/buyer.py index b4d8fb0..6d4d7a1 100644 --- a/generalresearch/models/thl/survey/buyer.py +++ b/generalresearch/models/thl/survey/buyer.py @@ -1,7 +1,9 @@ +from __future__ import annotations + from datetime import datetime, timezone from decimal import Decimal from math import log -from typing import Annotated, Optional +from typing import Annotated from pydantic import ( BaseModel, @@ -30,9 +32,9 @@ class Buyer(BaseModel): model_config = ConfigDict(validate_assignment=True) - id: Optional[PositiveInt] = Field(default=None, exclude=True) + id: PositiveInt | None = Field(default=None, exclude=True) # todo: need to add to db - uuid: Optional[UUIDStr] = Field(default=None) + uuid: UUIDStr | None = Field(default=None) source: Source = Field( description="The marketplace this buyer is on.\n" + Source.as_openapi() @@ -42,7 +44,7 @@ class Buyer(BaseModel): max_length=128, description="The internal code on this marketplace for this buyer", ) - label: Optional[str] = Field(default=None, max_length=255) + label: str | None = Field(default=None, max_length=255) created: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc), description="When this entry was made, or when the buyer was first seen", @@ -72,7 +74,7 @@ class BuyerWithDetail(BaseModel): """For API Responses""" buyer: Buyer = Field() - activity: Optional[BuyerActivity] = Field(default=None) + activity: BuyerActivity | None = Field(default=None) class BuyerCountryStat(BaseModel): @@ -89,12 +91,12 @@ class BuyerCountryStat(BaseModel): model_config = ConfigDict(validate_assignment=False) # ---- Identity ---- - buyer_id: Optional[PositiveInt] = Field( + buyer_id: PositiveInt | None = Field( default=None, exclude=True, description="This is the pk of the Buyer object in the db", ) - country_iso: Optional[CountryISOLike] = Field( + country_iso: CountryISOLike | None = Field( default=None, description="If null, this is a weighted average across all countries", examples=["us"], diff --git a/generalresearch/models/thl/survey/condition.py b/generalresearch/models/thl/survey/condition.py index 3610750..927b7e1 100644 --- a/generalresearch/models/thl/survey/condition.py +++ b/generalresearch/models/thl/survey/condition.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import hashlib from abc import ABC from enum import Enum from functools import cached_property -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any from pydantic import ( BaseModel, @@ -60,18 +62,18 @@ class MarketplaceCondition(BaseModel, ABC): negate: bool = Field(default=False) # ---- These fields should be overridden in the implementor --- - question_id: Optional[str] = Field(frozen=True) - values: List[str] = Field() + question_id: str | None = Field(frozen=True) + values: list[str] = Field() # These question_ids get converted to list value types - _CONVERT_LIST_TO_RANGE: List[str] = PrivateAttr(default_factory=list) + _CONVERT_LIST_TO_RANGE: list[str] = PrivateAttr(default_factory=list) @field_validator("values", mode="after") - def sort_values(cls, values: List[str]): + def sort_values(cls, values: list[str]): return sorted(values) @field_validator("values", mode="after") - def check_values_lower(cls, values: List[str]): + def check_values_lower(cls, values: list[str]): assert values == [s.lower() for s in values], "values must be lowercase" return values @@ -210,7 +212,7 @@ class MarketplaceCondition(BaseModel, ABC): return sum(len(v) for v in self.values) @cached_property - def values_ranges(self) -> List[Tuple[float, float]]: + def values_ranges(self) -> list[tuple[float, float]]: assert ( self.value_type == ConditionValueType.RANGE ), "only call this method when value_type is RANGE" @@ -227,11 +229,11 @@ class MarketplaceCondition(BaseModel, ABC): return hashlib.md5(s.encode()).hexdigest()[:7] @classmethod - def from_mysql(cls, d: Dict[str, Any]) -> Self: + def from_mysql(cls, d: dict[str, Any]) -> Self: d["values"] = d["values"][1:-1].split("|") if d["values"][1:-1] else [] return cls.model_validate(d) - def to_mysql(self) -> Dict[str, str]: + def to_mysql(self) -> dict[str, str]: # This is what is stored in the xxx_criterion table d = self.model_dump( mode="json", @@ -274,9 +276,9 @@ class MarketplaceCondition(BaseModel, ABC): def evaluate_criterion( self, - user_qas: Dict[str, Set[str]], - user_groups: Optional[Set[str]] = None, - ) -> Optional[bool]: + user_qas: dict[str, set[str]], + user_groups: set[str] | None = None, + ) -> bool | None: """Given this user's MRPQs, do they "pass" this criterion? :param user_qas: user's quals. Looks like {'qid1': {'ans1', 'ans2'}} diff --git a/generalresearch/models/thl/survey/model.py b/generalresearch/models/thl/survey/model.py index 8d37af4..3794c00 100644 --- a/generalresearch/models/thl/survey/model.py +++ b/generalresearch/models/thl/survey/model.py @@ -1,6 +1,8 @@ +from __future__ import annotations + from datetime import datetime, timezone from decimal import Decimal -from typing import Any, Dict, List, Optional, Tuple +from typing import Any from pydantic import ( BaseModel, @@ -33,7 +35,7 @@ class SurveyCategoryModel(BaseModel): model_config = ConfigDict(from_attributes=True) category: Category = Field() - strength: Optional[float] = Field(default=None) + strength: float | None = Field(default=None) class SurveyEligibilityDefinition(BaseModel): @@ -45,7 +47,7 @@ class SurveyEligibilityDefinition(BaseModel): """ # References a marketplace-specific question - property_codes: Tuple[PropertyCode, ...] = Field(default_factory=tuple) + property_codes: tuple[PropertyCode, ...] = Field(default_factory=tuple) @model_validator(mode="after") def sort_question_ids(self): @@ -56,16 +58,16 @@ class SurveyEligibilityDefinition(BaseModel): class Survey(BaseModel): model_config = ConfigDict(validate_assignment=True) - id: Optional[PositiveInt] = Field(default=None, exclude=True) + id: PositiveInt | None = Field(default=None, exclude=True) source: Source = Field() survey_id: str = Field(min_length=1, max_length=32, examples=["127492892"]) - buyer_id: Optional[int] = Field( + buyer_id: int | None = Field( default=None, exclude=True, description="This is the DB's fk id" ) # ---v So the fk id can be looked up from the code - buyer_code: Optional[str] = Field( + buyer_code: str | None = Field( min_length=1, max_length=128, default=None, examples=["124"] ) @@ -79,9 +81,9 @@ class Survey(BaseModel): is_live: bool = Field(default=True) is_recontact: bool = Field(default=False) - categories: List[SurveyCategoryModel] = Field(default_factory=list) + categories: list[SurveyCategoryModel] = Field(default_factory=list) - eligibility_criteria: Optional[SurveyEligibilityDefinition] = Field(default=None) + eligibility_criteria: SurveyEligibilityDefinition | None = Field(default=None) @property def natural_key(self) -> SurveyKey: @@ -119,10 +121,10 @@ class Survey(BaseModel): class SurveyStat(BaseModel): model_config = ConfigDict(validate_assignment=True) - id: Optional[PositiveInt] = Field(exclude=True, default=None) + id: PositiveInt | None = Field(exclude=True, default=None) # ---- Identity ---- - survey_id: Optional[PositiveInt] = Field( + survey_id: PositiveInt | None = Field( default=None, exclude=True, description="This is the pk of the Survey object in the db", @@ -137,8 +139,8 @@ class SurveyStat(BaseModel): # --- For lookup / de-normalization, to avoid potentially costly # joins on marketplace_survey table --- - survey_source: Optional[Source] = Field(default=None, exclude=True) - survey_survey_id: Optional[str] = Field( + survey_source: Source | None = Field(default=None, exclude=True) + survey_survey_id: str | None = Field( default=None, exclude=True, min_length=1, max_length=32 ) survey_is_live: bool = Field(default=True, exclude=True) @@ -201,7 +203,7 @@ class SurveyStat(BaseModel): return f"{self.survey_source.value}:{self.survey_survey_id}" @property - def unique_key(self) -> Tuple[int, Optional[str], str, int]: + def unique_key(self) -> tuple[int, str | None, str, int]: return self.survey_id, self.quota_id, self.country_iso, self.version def model_dump_sql(self): @@ -219,16 +221,16 @@ class TaskActivity(BaseModel): source: Source = Field() survey_id: str = Field(min_length=1, max_length=32, examples=["127492892"]) - status_counts: Dict[Status, NonNegativeInt] = Field(default_factory=dict) - status_code_1_counts: Dict[StatusCode1, NonNegativeInt] = Field( + status_counts: dict[Status, NonNegativeInt] = Field(default_factory=dict) + status_code_1_counts: dict[StatusCode1, NonNegativeInt] = Field( default_factory=dict ) in_progress_count: NonNegativeInt = Field( default=0, description="Count of entrances that have no Status and were entered within the past 90 minutes", ) - last_complete: Optional[AwareDatetimeISO] = Field(default=None) - last_entrance: Optional[AwareDatetimeISO] = Field(default=None) + last_complete: AwareDatetimeISO | None = Field(default=None) + last_entrance: AwareDatetimeISO | None = Field(default=None) @computed_field @property @@ -243,7 +245,7 @@ class TaskActivity(BaseModel): # ---- percentages ---- @computed_field @property - def status_percentages(self) -> Dict[Status, NonNegativeFloat]: + def status_percentages(self) -> dict[Status, NonNegativeFloat]: total = self.total_finished if total == 0: return {} @@ -251,7 +253,7 @@ class TaskActivity(BaseModel): @computed_field @property - def status_code_1_percentages(self) -> Dict[StatusCode1, NonNegativeFloat]: + def status_code_1_percentages(self) -> dict[StatusCode1, NonNegativeFloat]: total = sum(self.status_code_1_counts.values()) if total == 0: return {} @@ -259,21 +261,21 @@ class TaskActivity(BaseModel): class TaskActivityPublic(BaseModel): - source: Optional[Source] = Field(exclude=True, default=None) - survey_id: Optional[str] = Field( + source: Source | None = Field(exclude=True, default=None) + survey_id: str | None = Field( min_length=1, max_length=32, examples=["127492892"], exclude=True, default=None ) - status_percentages: Dict[Status, NonNegativeFloat] = Field(default_factory=dict) - status_code_1_percentages: Dict[ + status_percentages: dict[Status, NonNegativeFloat] = Field(default_factory=dict) + status_code_1_percentages: dict[ Annotated[StatusCode1, EnumNameSerializer], NonNegativeFloat ] = Field(default_factory=dict) - last_complete: Optional[AwareDatetimeISO] = Field(default=None) - last_entrance: Optional[AwareDatetimeISO] = Field(default=None) + last_complete: AwareDatetimeISO | None = Field(default=None) + last_entrance: AwareDatetimeISO | None = Field(default=None) @field_validator("status_code_1_percentages", mode="before") - def transform_enum_name_pct(cls, value: Dict[str, Any]) -> Dict[str, Any]: + def transform_enum_name_pct(cls, value: dict[str, Any]) -> dict[str, Any]: # If we are serializing+deserializing this model (i.e. when we cache # it), this fails because we've replaced the enum value with the # name. Put it back here ... @@ -287,8 +289,8 @@ class TaskActivityPublic(BaseModel): class TaskActivityPrivate(TaskActivityPublic): - status_counts: Dict[Status, int] = Field(default_factory=dict) - status_code_1_counts: Dict[Annotated[StatusCode1, EnumNameSerializer], int] = Field( + status_counts: dict[Status, int] = Field(default_factory=dict) + status_code_1_counts: dict[Annotated[StatusCode1, EnumNameSerializer], int] = Field( default_factory=dict ) in_progress_count: NonNegativeInt = Field( @@ -297,7 +299,7 @@ class TaskActivityPrivate(TaskActivityPublic): ) @field_validator("status_code_1_counts", mode="before") - def transform_enum_name_cnt(cls, value: Dict[str, Any]) -> Dict[str, Any]: + def transform_enum_name_cnt(cls, value: dict[str, Any]) -> dict[str, Any]: # If we are serializing+deserializing this model (i.e. when we cache # it), this fails because we've replaced the enum value with the # name. Put it back here ... @@ -310,12 +312,12 @@ class TaskWithDetail(BaseModel): """For API Responses""" task: Survey = Field() - stats: List[SurveyStat] = Field(default_factory=list) - activity_global: Optional[TaskActivityPublic] = Field(default=None) - activity_product: Optional[TaskActivityPrivate] = Field(default=None) + stats: list[SurveyStat] = Field(default_factory=list) + activity_global: TaskActivityPublic | None = Field(default=None) + activity_product: TaskActivityPrivate | None = Field(default=None) class TasksWithDetail(Page): """For API Responses""" - tasks: List[TaskWithDetail] = Field(default_factory=list) + tasks: list[TaskWithDetail] = Field(default_factory=list) diff --git a/generalresearch/models/thl/survey/penalty.py b/generalresearch/models/thl/survey/penalty.py index a9c8a56..e9515d4 100644 --- a/generalresearch/models/thl/survey/penalty.py +++ b/generalresearch/models/thl/survey/penalty.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import abc from datetime import datetime, timezone -from typing import List, Literal, Union +from typing import Literal from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from typing_extensions import Annotated @@ -60,4 +62,4 @@ Penalty = Annotated[ Union[BPSurveyPenalty, TeamSurveyPenalty], Field(discriminator="kind"), ] -PenaltyListAdapter = TypeAdapter(List[Penalty]) +PenaltyListAdapter = TypeAdapter(list[Penalty]) diff --git a/generalresearch/models/thl/survey/task_collection.py b/generalresearch/models/thl/survey/task_collection.py index 804fb49..b09cdfb 100644 --- a/generalresearch/models/thl/survey/task_collection.py +++ b/generalresearch/models/thl/survey/task_collection.py @@ -1,7 +1,8 @@ +from __future__ import annotations + import copy import json import logging -from typing import List import pandas as pd from pandera import DataFrameSchema @@ -24,7 +25,7 @@ class TaskCollection(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) # overload this with the correct type! - items: List[MarketplaceTask] + items: list[MarketplaceTask] df: pd.DataFrame = Field(default_factory=pd.DataFrame) # overload this with the correct schema! diff --git a/generalresearch/models/thl/synchronize_global_vars.py b/generalresearch/models/thl/synchronize_global_vars.py index 72d987d..a656088 100644 --- a/generalresearch/models/thl/synchronize_global_vars.py +++ b/generalresearch/models/thl/synchronize_global_vars.py @@ -1,4 +1,4 @@ -from typing import List +from __future__ import annotations from pydantic import BaseModel, Field @@ -10,7 +10,7 @@ class SynchronizeGlobalVarsMsg(BaseModel): # Specifies the key / subkey to be acted upon. # For example ["mrpq", 123] would apply to GLOBAL_VARS["mrpq"][123] - key_path: List[str] = Field() + key_path: list[str] = Field() # e.g. GLOBAL_VARS["mrpq"].pop(123, None) pop: bool = Field(default=False) # e.g. GLOBAL_VARS["mrpq"][123].clear() diff --git a/generalresearch/models/thl/task_adjustment.py b/generalresearch/models/thl/task_adjustment.py index 9d8a04d..89a3873 100644 --- a/generalresearch/models/thl/task_adjustment.py +++ b/generalresearch/models/thl/task_adjustment.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from datetime import datetime, timezone from decimal import Decimal -from typing import Optional from uuid import uuid4 from pydantic import BaseModel, ConfigDict, Field, PositiveInt, model_validator @@ -55,8 +56,8 @@ class TaskAdjustmentEvent(BaseModel): # Same thing as with adjusted_status, the amount is the "amount the cpi # is changing by"!! - amount: Optional[Decimal] = Field(lt=1000, ge=-1000, default=None) - ext_status_code: Optional[str] = Field(default=None, max_length=32) + amount: Decimal | None = Field(lt=1000, ge=-1000, default=None) + ext_status_code: str | None = Field(default=None, max_length=32) wall_uuid: UUIDStr = Field(description="The wall event being adjusted") diff --git a/generalresearch/models/thl/task_status.py b/generalresearch/models/thl/task_status.py index a7aec91..ee713a5 100644 --- a/generalresearch/models/thl/task_status.py +++ b/generalresearch/models/thl/task_status.py @@ -1,5 +1,7 @@ +from __future__ import annotations + from datetime import datetime -from typing import Annotated, Any, Dict, List, Literal, Optional +from typing import Annotated, Any, Literal from pydantic import ( BaseModel, @@ -69,12 +71,12 @@ class TaskStatusResponse(BaseModel): started: AwareDatetimeISO = Field(description="When the session was started") - finished: Optional[AwareDatetimeISO] = Field( + finished: AwareDatetimeISO | None = Field( default=None, description="When the session was finished" ) # This uses the grpc's Status enum. It gets serialized to an int. - status: Optional[Status] = Field( + status: Status | None = Field( default=None, examples=[3], description="The outcome of a session." @@ -84,14 +86,14 @@ class TaskStatusResponse(BaseModel): " - 3 - COMPLETE (the user completed the task)", ) - payout: Optional[NonNegativeInt] = Field( + payout: NonNegativeInt | None = Field( default=None, lt=100_000, examples=[500], description="The amount paid to the supplier, in integer USD cents", ) - user_payout: Optional[NonNegativeInt] = Field( + user_payout: NonNegativeInt | None = Field( default=None, lt=100_000, description="If a payout transformation is configured on this account, " @@ -99,78 +101,78 @@ class TaskStatusResponse(BaseModel): examples=[337], ) - payout_format: Optional[PayoutFormatType] = PayoutFormatOptionalField + payout_format: PayoutFormatType | None = PayoutFormatOptionalField - user_payout_string: Optional[str] = Field( + user_payout_string: str | None = Field( default=None, description="If a payout transformation is configured on this account, " "this is the amount to display to the user", examples=["3370 Points"], ) - kwargs: Dict[str, str] = Field( + kwargs: dict[str, str] = Field( default_factory=dict, description="Any extra url params used in the offerwall request will be " "passed back here", ) - status_code_1: Optional[Annotated[StatusCode1, EnumNameSerializer]] = Field( + status_code_1: Annotated[StatusCode1, EnumNameSerializer] | None = Field( default=None, examples=[StatusCode1.COMPLETE.name], description=StatusCode1.as_openapi_with_value_descriptions_name(), ) - status_code_2: Optional[Annotated[SessionStatusCode2, EnumNameSerializer]] = Field( + status_code_2: Annotated[SessionStatusCode2, EnumNameSerializer] | None = Field( default=None, examples=[None], description=SessionStatusCode2.as_openapi_with_value_descriptions_name(), ) - adjusted_status: Optional[SessionAdjustedStatus] = Field( + adjusted_status: SessionAdjustedStatus | None = Field( default=None, description=SessionAdjustedStatus.as_openapi_with_value_descriptions(), examples=[None], ) - adjusted_timestamp: Optional[AwareDatetimeISO] = Field( + adjusted_timestamp: AwareDatetimeISO | None = Field( default=None, description="When the adjusted status was last set.", examples=[None], ) - adjusted_payout: Optional[NonNegativeInt] = Field( + adjusted_payout: NonNegativeInt | None = Field( default=None, lt=100_000, description="The new payout after adjustment.", examples=[None], ) - adjusted_user_payout: Optional[NonNegativeInt] = Field( + adjusted_user_payout: NonNegativeInt | None = Field( default=None, lt=100_000, description="The new user_payout after adjustment.", examples=[None], ) - adjusted_user_payout_string: Optional[str] = Field( + adjusted_user_payout_string: str | None = Field( default=None, description="The new user_payout_string after adjustment.", examples=[None], ) # This is used for validation purposes only. It won't get serialized - payout_transformation: Optional[PayoutTransformation] = Field( + payout_transformation: PayoutTransformation | None = Field( default=None, exclude=True ) - wall_events: Optional[List[WallOut]] = Field(default=None) + wall_events: list[WallOut] | None = Field(default=None) currency: Literal["USD"] = Field(default="USD") final_status: int = Field(default=0, description="This is deprecated") # Serialize enum → int @field_serializer("status", return_type=int) - def serialize_status(self, v: Optional[Status], _info): + def serialize_status(self, v: Status | None, _info): return STATUS_MAP[v] # Accept int OR string for input, but internally store a Status enum @@ -219,11 +221,11 @@ class TaskStatusResponse(BaseModel): return v @field_validator("payout", mode="before") - def transform_payout(cls, v: Optional[NonNegativeInt]) -> NonNegativeInt: + def transform_payout(cls, v: NonNegativeInt | None) -> NonNegativeInt: return v or 0 @field_validator("kwargs", mode="after") - def sanitize_kwargs(cls, v: Optional[Dict]) -> Optional[Dict]: + def sanitize_kwargs(cls, v: dict | None) -> dict | None: if v and "clicked_timestamp" in v: try: clicked_timestamp = datetime.strptime( @@ -289,4 +291,4 @@ class TaskStatusResponse(BaseModel): class TasksStatusResponse(Page): - tasks_status: List[TaskStatusResponse] = Field(default_factory=list) + tasks_status: list[TaskStatusResponse] = Field(default_factory=list) diff --git a/generalresearch/models/thl/user.py b/generalresearch/models/thl/user.py index 852766c..34c9b35 100644 --- a/generalresearch/models/thl/user.py +++ b/generalresearch/models/thl/user.py @@ -4,7 +4,7 @@ import json import logging import re from datetime import datetime, timezone -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING from uuid import UUID, uuid4 from pydantic import ( @@ -45,22 +45,22 @@ BPUID_ALLOWED = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ class User(BaseModel): model_config = ConfigDict(extra="forbid", validate_assignment=True) - user_id: Optional[PositiveInt] = Field( + user_id: PositiveInt | None = Field( default=None, lt=MAX_INT32, serialization_alias="id" ) - uuid: Optional[UUIDStr] = Field(default=None, examples=[uuid4().hex]) + uuid: UUIDStr | None = Field(default=None, examples=[uuid4().hex]) # 'product' is a Class with values that are fetched from the DB. # Initialization is deferred until it is actually needed # (see .prefetch_product()) - product: Optional[Product] = Field(default=None) + product: Product | None = Field(default=None) - product_id: Optional[UUIDStr] = Field( + product_id: UUIDStr | None = Field( default=None, examples=["4fe381fb7186416cb443a38fa66c6557"] ) - product_user_id: Optional[BPUIDStr] = Field( + product_user_id: BPUIDStr | None = Field( default=None, examples=["app-user-9329ebd"], description="A unique identifier for each user, which is set by the " @@ -72,24 +72,24 @@ class User(BaseModel): # TODO: Is it possible to protect these from ever being initialized? # - Would need to be allowed with .from_json but not User constructor directly # - Would need to allow private setters for setting from DB values - blocked: Optional[bool] = Field(default=False, strict=True) + blocked: bool | None = Field(default=False, strict=True) - created: Optional[AwareDatetimeISO] = Field( + created: AwareDatetimeISO | None = Field( default=None, description="When the user was created on the GRL platform.", ) # Note: due to cacheing, last_seen might be up to a day out of date! - last_seen: Optional[AwareDatetimeISO] = Field( + last_seen: AwareDatetimeISO | None = Field( default=None, description="When the user was last seen on, or acting on any" "part of the GRL platform.", ) # --- Prefetch Fields --- - audit_log: Optional[List[AuditLog]] = Field(default=None) - transactions: Optional[List["LedgerTransaction"]] = Field(default=None) - location_history: Optional[List["GeoIPInformation"]] = Field(default=None) + audit_log: list[AuditLog] | None = Field(default=None) + transactions: list["LedgerTransaction"] | None = Field(default=None) + location_history: list["GeoIPInformation"] | None = Field(default=None) # --- Prebuild Fields --- # session: Optional[List] = Field(default=None) @@ -199,7 +199,7 @@ class User(BaseModel): raise ValueError("product_user_id must not equal the product_id") return True - def to_dict(self) -> Dict: + def to_dict(self) -> dict: return self.model_dump(mode="python", exclude={"product"}) def to_json(self) -> str: @@ -279,16 +279,12 @@ class User(BaseModel): pm = ProductManager(pg_config=pg_config) self.product = pm.get_by_uuid(product_uuid=self.product_id) - return None - def prefetch_audit_log(self, audit_log_manager: "AuditLogManager") -> None: self.audit_log = audit_log_manager.filter_by_user_id(user_id=self.user_id) - return None def prefetch_transactions(self, thl_lm: "ThlLedgerManager") -> None: account = thl_lm.get_account_or_create_user_wallet(user=self) self.transactions = thl_lm.get_tx_filtered_by_account(account_uuid=account.uuid) - return None # def prefetch_location_history(self, user_ip_history_manager: "UserIpHistoryManager") -> None: # return user_ip_history_manager.get_user_ip_history(user_id=self.user_id) diff --git a/generalresearch/models/thl/user_iphistory.py b/generalresearch/models/thl/user_iphistory.py index 159b97b..84ee42e 100644 --- a/generalresearch/models/thl/user_iphistory.py +++ b/generalresearch/models/thl/user_iphistory.py @@ -1,6 +1,7 @@ +from __future__ import annotations + import ipaddress from datetime import datetime, timedelta, timezone -from typing import Dict, List, Optional from faker import Faker from pydantic import ( @@ -34,10 +35,10 @@ class UserIPRecord(BaseModel): ip: IPvAnyAddressStr = Field() created: AwareDatetimeISO = Field() - information: Optional[GeoIPInformation] = Field(default=None, exclude=True) + information: GeoIPInformation | None = Field(default=None, exclude=True) @property - def country_iso(self) -> Optional[CountryISOLike]: + def country_iso(self) -> CountryISOLike | None: return self.information.country_iso if self.information else None @property @@ -52,15 +53,15 @@ class UserIPRecord(BaseModel): ) @property - def user_type(self) -> Optional[UserType]: + def user_type(self) -> UserType | None: return self.information.user_type if self.information else None @property - def subdivision_1_iso(self) -> Optional[str]: + def subdivision_1_iso(self) -> str | None: return self.information.subdivision_1_iso if self.information else None @property - def subdivision_2_iso(self) -> Optional[str]: + def subdivision_2_iso(self) -> str | None: return self.information.subdivision_2_iso if self.information else None @@ -71,14 +72,12 @@ class IPRecord(BaseModel): # On a top-level, this should be an empty list if there are no forwarded_ip. # Within a forwarded_ip record, this should be None. - forwarded_ip_records: Optional[List["IPRecord"]] = Field( - default=None, description="" - ) + forwarded_ip_records: list["IPRecord"] | None = Field(default=None, description="") - information: Optional[GeoIPInformation] = Field(default=None) + information: GeoIPInformation | None = Field(default=None) @property - def forwarded_ips(self) -> Optional[List[IPvAnyAddressStr]]: + def forwarded_ips(self) -> list[IPvAnyAddressStr] | None: return ( [x.ip for x in self.forwarded_ip_records] if self.forwarded_ip_records is not None @@ -86,7 +85,7 @@ class IPRecord(BaseModel): ) def ip_changed( - self, ip: IPvAnyAddressStr, forwarded_ips: List[IPvAnyAddressStr] + self, ip: IPvAnyAddressStr, forwarded_ips: list[IPvAnyAddressStr] ) -> bool: return not (ip == self.ip and forwarded_ips == self.forwarded_ips) @@ -110,11 +109,10 @@ class IPRecord(BaseModel): x.information = res.get(x.ip) else: self.information = m.get(ip_address=self.ip) - return None # --- ORM --- @classmethod - def from_mysql(cls, d: Dict) -> Self: + def from_mysql(cls, d: dict) -> Self: created = d["created"].replace(tzinfo=timezone.utc) d["created"] = created @@ -149,22 +147,22 @@ class UserIPHistory(BaseModel): # In thl-gprc, we run "audit_ip_history()", and so a user should # get blocked after 100 IP switches or 30 unique IPs # Sorted created DESC - ips: Optional[List[UserIPRecord]] = Field( + ips: list[UserIPRecord] | None = Field( default=None, description="These are any IP addresses that came in ", max_length=101, ) - ips_ws: Optional[List[IPRecord]] = Field( + ips_ws: list[IPRecord] | None = Field( default=None, description="These are any IP addresses that came in " ) - ips_dns: Optional[List[IPRecord]] = Field( + ips_dns: list[IPRecord] | None = Field( default=None, description="These are any IP addresses that came in " ) # -- prefetch_ fields - user: Optional[User] = Field(default=None) + user: User | None = Field(default=None) @field_validator("ips", mode="after") @classmethod @@ -195,8 +193,6 @@ class UserIPHistory(BaseModel): ) self.user = um.get_user(user_id=self.user_id) - return None - def enrich_ips(self, pg_config: PostgresConfig, redis_config: RedisConfig) -> None: from generalresearch.managers.thl.ipinfo import GeoIpInfoManager diff --git a/generalresearch/models/thl/user_profile.py b/generalresearch/models/thl/user_profile.py index 98b3326..e96266a 100644 --- a/generalresearch/models/thl/user_profile.py +++ b/generalresearch/models/thl/user_profile.py @@ -1,5 +1,7 @@ +from __future__ import annotations + import hashlib -from typing import Any, Dict, List, Optional +from typing import Any from pydantic import ( BaseModel, @@ -21,19 +23,17 @@ from generalresearch.models.thl.user_streak import UserStreak class UserMetadata(BaseModel): model_config = ConfigDict(extra="forbid", validate_assignment=True) - user_id: SkipJsonSchema[Optional[PositiveInt]] = Field( + user_id: SkipJsonSchema[PositiveInt | None] = Field( exclude=True, default=None, lt=MAX_INT32 ) - email_address: Optional[EmailStr] = Field( - default=None, examples=["contact@mail.com"] - ) + email_address: EmailStr | None = Field(default=None, examples=["contact@mail.com"]) @computed_field def email_md5( self, ) -> Annotated[ - Optional[str], + str | None, Field( min_length=32, max_length=32, @@ -50,7 +50,7 @@ class UserMetadata(BaseModel): def email_sha1( self, ) -> Annotated[ - Optional[str], + str | None, Field( min_length=40, max_length=40, @@ -66,7 +66,7 @@ class UserMetadata(BaseModel): def email_sha256( self, ) -> Annotated[ - Optional[str], + str | None, Field( min_length=64, max_length=64, @@ -80,7 +80,7 @@ class UserMetadata(BaseModel): return None return hashlib.sha256(self.email_address.encode("utf-8")).hexdigest() - def to_db(self) -> Dict[str, Any]: + def to_db(self) -> dict[str, Any]: res = self.model_dump(mode="json") res["user_id"] = self.user_id return res @@ -107,7 +107,7 @@ class UserProfile(UserMetadata): user: User = Field() - marketplace_pids: Dict[Source, UUIDStr] = Field( + marketplace_pids: dict[Source, UUIDStr] = Field( default_factory=dict, description="User's PID in marketplaces", examples=[ @@ -119,4 +119,4 @@ class UserProfile(UserMetadata): ], ) - streaks: List[UserStreak] = Field(default_factory=list) + streaks: list[UserStreak] = Field(default_factory=list) diff --git a/generalresearch/models/thl/user_quality_event.py b/generalresearch/models/thl/user_quality_event.py index 2b4873a..d6ebddc 100644 --- a/generalresearch/models/thl/user_quality_event.py +++ b/generalresearch/models/thl/user_quality_event.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime, timezone from decimal import Decimal from enum import Enum -from typing import List, Literal, Optional +from typing import Literal from pydantic import BaseModel, Field, PositiveInt @@ -46,15 +46,15 @@ class AccessControlEvent(BaseModel): QualityEventType.clear_access_control_list, ] = Field() # One of user_id / (product_id, bpuid) is required. - product_id: Optional[UUIDStr] = Field( + product_id: UUIDStr | None = Field( default=None, examples=["4fe381fb7186416cb443a38fa66c6557"] ) - bpuid: Optional[BPUIDStr] = Field(default=None, examples=["app-user-9329ebd"]) - user_id: Optional[PositiveInt] = Field(default=None, lt=MAX_INT32) + bpuid: BPUIDStr | None = Field(default=None, examples=["app-user-9329ebd"]) + user_id: PositiveInt | None = Field(default=None, lt=MAX_INT32) class AccessControlEventBody(BaseModel): - events: List[AccessControlEvent] = Field(max_length=100, min_length=1) + events: list[AccessControlEvent] = Field(max_length=100, min_length=1) class TaskAdjustmentEvent(BaseModel): @@ -70,21 +70,21 @@ class TaskAdjustmentEvent(BaseModel): # Only MID is needed to populate all the following, however we can pass them in order # to perform validation. If any disagree, an error should be raised. - survey_id: Optional[str] = Field(max_length=32, default=None) - amount: Optional[Decimal] = Field( + survey_id: str | None = Field(max_length=32, default=None) + amount: Decimal | None = Field( description="If negative, the status should adjusted to incomplete", default=None, ) - event_time: Optional[AwareDatetimeISO] = Field( + event_time: AwareDatetimeISO | None = Field( description="This is when the original wall event was started", default=None, ) - product_id: Optional[UUIDStr] = Field( + product_id: UUIDStr | None = Field( default=None, examples=["4fe381fb7186416cb443a38fa66c6557"] ) - bpuid: Optional[BPUIDStr] = Field(default=None, examples=["app-user-9329ebd"]) - user_id: Optional[PositiveInt] = Field(default=None, lt=MAX_INT32) + bpuid: BPUIDStr | None = Field(default=None, examples=["app-user-9329ebd"]) + user_id: PositiveInt | None = Field(default=None, lt=MAX_INT32) class TaskAdjustmentEventBody(BaseModel): - events: List[TaskAdjustmentEvent] = Field(max_length=100, min_length=1) + events: list[TaskAdjustmentEvent] = Field(max_length=100, min_length=1) diff --git a/generalresearch/models/thl/user_streak.py b/generalresearch/models/thl/user_streak.py index 5d13bc5..278809b 100644 --- a/generalresearch/models/thl/user_streak.py +++ b/generalresearch/models/thl/user_streak.py @@ -1,6 +1,7 @@ +from __future__ import annotations + from datetime import date, datetime, timedelta from enum import Enum -from typing import Optional, Tuple import pandas as pd from pydantic import ( @@ -63,7 +64,7 @@ class UserStreak(BaseModel): ser_json_timedelta="float", validate_assignment=True, extra="forbid" ) - user_id: SkipJsonSchema[Optional[PositiveInt]] = Field( + user_id: SkipJsonSchema[PositiveInt | None] = Field( exclude=True, default=None, lt=MAX_INT32 ) country_iso: CountryISO = Field() @@ -75,7 +76,7 @@ class UserStreak(BaseModel): current_streak: NonNegativeInt = Field() longest_streak: NonNegativeInt = Field() state: StreakState = Field() - last_fulfilled_period_start: Optional[date] = Field(default=None) + last_fulfilled_period_start: date | None = Field(default=None) @computed_field() @property @@ -92,17 +93,17 @@ class UserStreak(BaseModel): @computed_field() @property - def current_period_bounds(self) -> Tuple[AwareDatetime, AwareDatetime]: + def current_period_bounds(self) -> tuple[AwareDatetime, AwareDatetime]: return self.get_period_bounds(datetime.now(tz=self.timezone).date()) @computed_field() @property - def last_fulfilled_period_bounds(self) -> Optional[Tuple[datetime, datetime]]: + def last_fulfilled_period_bounds(self) -> tuple[datetime, datetime] | None: return self.get_period_bounds(self.last_fulfilled_period_start) @computed_field() @property - def time_remaining_in_period(self) -> Optional[timedelta]: + def time_remaining_in_period(self) -> timedelta | None: # Time left to continue your streak if self.state in {StreakState.BROKEN, StreakState.ACTIVE}: return None @@ -129,9 +130,7 @@ class UserStreak(BaseModel): ), "Current streak can't be longer than longest streak" return self - def get_period_bounds( - self, start_date: date - ) -> Optional[Tuple[datetime, datetime]]: + def get_period_bounds(self, start_date: date) -> tuple[datetime, datetime] | None: """ Returns (period_start_local, period_end_local) Both timezone-aware. diff --git a/generalresearch/models/thl/userhealth.py b/generalresearch/models/thl/userhealth.py index fd275de..e556dc8 100644 --- a/generalresearch/models/thl/userhealth.py +++ b/generalresearch/models/thl/userhealth.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import datetime, timezone from enum import Enum from typing import Dict, Optional diff --git a/generalresearch/models/thl/wallet/cashout_method.py b/generalresearch/models/thl/wallet/cashout_method.py index c04bd66..59cf721 100644 --- a/generalresearch/models/thl/wallet/cashout_method.py +++ b/generalresearch/models/thl/wallet/cashout_method.py @@ -4,7 +4,7 @@ import hashlib import logging from datetime import datetime, timezone from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Literal from pydantic import ( BaseModel, @@ -62,22 +62,22 @@ class CashoutMethodBase(BaseModel): default="USD", description="The currency of the cashout. Only USD is supported.", ) - original_currency: Optional[Currency] = Field( + original_currency: Currency | None = Field( default=None, description="The base currency of the money paid out. This is used for " "e.g. sending an Amazon UK gift card", ) # This also is used for the PayoutEvent.request_data - data: Union[ - PaypalCashoutMethodData, - TangoCashoutMethodData, - CashMailCashoutMethodData, - AmtCashoutMethodData, - ] = Field(discriminator="type") + data: ( + PaypalCashoutMethodData + | TangoCashoutMethodData + | CashMailCashoutMethodData + | AmtCashoutMethodData + ) = Field(discriminator="type") description: str = Field( description="The description of the cashout method.", default="" ) - image_url: Optional[HttpsUrlStr] = Field( + image_url: HttpsUrlStr | None = Field( description="Link to an image to display", default=None ) max_value: PositiveInt = Field( @@ -93,18 +93,18 @@ class CashoutMethodBase(BaseModel): type: PayoutType = Field( description=PayoutType.as_openapi_with_value_descriptions(), ) - ext_id: Optional[str] = Field( + ext_id: str | None = Field( default=None, description="An external ID. Can be shown to a user to disambiguate " "a user's possibly multiple methods", ) - usd_exchange_rate: Optional[float] = Field(default=None) - max_value_usd: Optional[USDCent] = Field( + usd_exchange_rate: float | None = Field(default=None) + max_value_usd: USDCent | None = Field( default=None, description="(In lowest unit of USD), " "The maximum amount that can be cashed out in one transaction.", ) - min_value_usd: Optional[USDCent] = Field( + min_value_usd: USDCent | None = Field( default=None, description="(In lowest unit of USD), " "The minimum amount that can be cashed out in one transaction.", @@ -138,7 +138,7 @@ class CashoutMethodBase(BaseModel): class CashoutMethod(CashoutMethodBase): - user: Optional[User] = Field( + user: User | None = Field( default=None, description="If set, this cashout method is custom for this user. For example" "a user may have a paypal cashout method with their paypal" @@ -163,11 +163,11 @@ class CashoutMethod(CashoutMethodBase): class CashoutMethodOut(CashoutMethodBase): - product_id: Optional[UUIDStr] = Field( + product_id: UUIDStr | None = Field( default=None, examples=["4fe381fb7186416cb443a38fa66c6557"] ) - product_user_id: Optional[str] = Field( + product_user_id: str | None = Field( default=None, min_length=3, max_length=128, @@ -189,12 +189,12 @@ class CashoutMethodOut(CashoutMethodBase): class USDeliveryAddress(BaseModel): name_or_attn: str = Field(min_length=1, max_length=50) - company: Optional[str] = Field( + company: str | None = Field( default=None, min_length=1, max_length=50, ) - phone_number: Optional[str] = Field( + phone_number: str | None = Field( default=None, min_length=10, max_length=10, @@ -231,13 +231,13 @@ class TangoCashoutMethodData(BaseModel): type: Literal[PayoutType.TANGO] = Field(default=PayoutType.TANGO) utid: str = Field(description="tango utid") # TODO: Can't be CountryISOLike because it appears to be allcaps - countries: List[str] = Field() + countries: list[str] = Field() value_type: Literal["variable", "fixed"] = Field() disclaimer: str = Field(default="") terms: str = Field(default="") @field_validator("countries", mode="after") - def countries_case(cls, countries: List[str]) -> List[str]: + def countries_case(cls, countries: list[str]) -> list[str]: return [x.lower() for x in countries] @@ -246,7 +246,7 @@ class AmtCashoutMethodData(BaseModel): class CashoutMethodsResponse(StatusResponse): - cashout_methods: List[CashoutMethodOut] = Field() + cashout_methods: list[CashoutMethodOut] = Field() class DeliveryStatus(str, Enum): @@ -279,37 +279,37 @@ class ShippingMethod(str, Enum): # This goes in the PayoutEvent.order_data class CashMailOrderData(BaseModel): type: Literal[PayoutType.CASH_IN_MAIL] = Field(default=PayoutType.CASH_IN_MAIL) - shipping_cost: Optional[PositiveInt] = Field( + shipping_cost: PositiveInt | None = Field( description="(USD cents) The shipping cost. This amount get charged to the BP.", strict=True, ) - tracking_number: Optional[str] = Field( + tracking_number: str | None = Field( default=None, min_length=1, max_length=50, ) - shipping_method: Optional[ShippingMethod] = Field( + shipping_method: ShippingMethod | None = Field( default=None, min_length=1, max_length=50, description="Standard, express, etc.", ) - carrier: Optional[ShippingCarrier] = Field( + carrier: ShippingCarrier | None = Field( default=None, min_length=1, max_length=50, description="Name of the shipping company, e.g., USPS, FedEx, DHL", ) - ship_date: Optional[AwareDatetimeISO] = Field(default=None) - estimated_delivery_date: Optional[AwareDatetimeISO] = Field(default=None) - delivery_status: Optional[DeliveryStatus] = Field( + ship_date: AwareDatetimeISO | None = Field(default=None) + estimated_delivery_date: AwareDatetimeISO | None = Field(default=None) + delivery_status: DeliveryStatus | None = Field( default=None, min_length=1, max_length=50, description="Current status of delivery, e.g., pending, in " "transit, delivered", ) - last_updated: Optional[AwareDatetimeISO] = Field( + last_updated: AwareDatetimeISO | None = Field( default=None, description="Timestamp of the last status update", ) @@ -364,7 +364,7 @@ class CashoutRequestInfo(BaseModel): Payout. This is used only in the API response. """ - id: Optional[UUIDStr] = Field( + id: UUIDStr | None = Field( description="Unique ID for this cashout. This may be NULL if the " "status is REJECTED or FAILED, which may happen if the " "request is invalid.", @@ -374,13 +374,13 @@ class CashoutRequestInfo(BaseModel): description="This is the name of the cashout method.", examples=["Visa® Prepaid Card USD"], ) - message: Optional[str] = Field(default=None) - status: Optional[PayoutStatus] = Field( + message: str | None = Field(default=None) + status: PayoutStatus | None = Field( default=PayoutStatus.PENDING, description=PayoutStatus.as_openapi(), examples=[PayoutStatus.PENDING], ) - transaction_info: Optional[Dict[str, Any]] = Field(default=None) + transaction_info: dict[str, Any] | None = Field(default=None) class CashoutRequestResponse(StatusResponse): diff --git a/generalresearch/models/thl/wallet/payout.py b/generalresearch/models/thl/wallet/payout.py index 72aea61..4b1f5a2 100644 --- a/generalresearch/models/thl/wallet/payout.py +++ b/generalresearch/models/thl/wallet/payout.py @@ -1,6 +1,9 @@ +from __future__ import annotations + import json +from collections.abc import Collection from datetime import datetime, timezone -from typing import Any, Collection, Dict, List, Optional, Union +from typing import Any from uuid import uuid4 from pydantic import ( @@ -35,8 +38,8 @@ class PayoutEvent(BaseModel, validate_assignment=True): # These two fields are copied here from the LedgerAccount through the # debit_account_uuid for convenience. They will get populated if the # PayoutEventManager retrieves a PayoutEvent from the db. - account_reference_type: Optional[str] = Field(default=None) - account_reference_uuid: Optional[UUIDStr] = Field(default=None) + account_reference_type: str | None = Field(default=None) + account_reference_uuid: UUIDStr | None = Field(default=None) # References a row in the account_cashoutmethod table. This is the # cashout method that was used to request this payout. (A cashout is @@ -46,7 +49,7 @@ class PayoutEvent(BaseModel, validate_assignment=True): # By default, this will just be the cashout_method.name. This also is # populated from the db and so does not need to be set (there is no # `description` field in event_payout) - description: Optional[str] = Field(default=None) + description: str | None = Field(default=None) created: AwareDatetimeISO = Field( default_factory=lambda: datetime.now(tz=timezone.utc) ) @@ -67,18 +70,18 @@ class PayoutEvent(BaseModel, validate_assignment=True): ) # Used for holding an external, payout-type-specific identifier - ext_ref_id: Optional[str] = Field(default=None) + ext_ref_id: str | None = Field(default=None) payout_type: PayoutType = Field( description=PayoutType.as_openapi(), examples=[PayoutType.ACH] ) # Stores payout-type-specific information that is used to request this # payout from the external provider. - request_data: Dict[str, Any] = Field(default_factory=dict) + request_data: dict[str, Any] = Field(default_factory=dict) # Stores payout-type-specific order information that is returned from # the external payout provider. - order_data: Optional[Union[Dict[str, Any], CashMailOrderData]] = Field(default=None) + order_data: dict[str, Any] | CashMailOrderData | None = Field(default=None) @field_validator("payout_type", mode="before") @classmethod @@ -93,8 +96,8 @@ class PayoutEvent(BaseModel, validate_assignment=True): def update( self, status: PayoutStatus, - ext_ref_id: Optional[str] = None, - order_data: Optional[Dict[str, Any]] = None, + ext_ref_id: str | None = None, + order_data: dict[str, Any] | None = None, ) -> None: # These 3 things are the only modifiable attributes self.check_status_change_allowed(status) @@ -128,7 +131,7 @@ class PayoutEvent(BaseModel, validate_assignment=True): else: raise ValueError("this shouldn't happen") - def model_dump_mysql(self, *args, **kwargs) -> Dict[str, Any]: + def model_dump_mysql(self, *args, **kwargs) -> dict[str, Any]: d = self.model_dump(mode="json", *args, **kwargs) if "created" in d: d["created"] = self.created.replace(tzinfo=None) @@ -166,7 +169,7 @@ class BPPayoutEvent(BaseModel): examples=[531], ) - status: Optional[PayoutStatus] = Field( + status: PayoutStatus | None = Field( default=PayoutStatus.PENDING, description=PayoutStatus.as_openapi(), examples=[PayoutStatus.COMPLETE], @@ -186,9 +189,9 @@ class BPPayoutEvent(BaseModel): @staticmethod def from_pe( payout_events: Collection[PayoutEvent], - account_product_mapping: Dict[str, str], - order_by="ASC", - ) -> List["BPPayoutEvent"]: + account_product_mapping: dict[str, str], + order_by: str = "ASC", + ) -> list["BPPayoutEvent"]: res = [] for pe in payout_events: bp_pe = BPPayoutEvent.model_validate( |