diff options
Diffstat (limited to 'generalresearch/managers/thl/profiling')
| -rw-r--r-- | generalresearch/managers/thl/profiling/question.py | 26 | ||||
| -rw-r--r-- | generalresearch/managers/thl/profiling/schema.py | 7 | ||||
| -rw-r--r-- | generalresearch/managers/thl/profiling/uqa.py | 30 | ||||
| -rw-r--r-- | generalresearch/managers/thl/profiling/user_upk.py | 34 |
4 files changed, 49 insertions, 48 deletions
diff --git a/generalresearch/managers/thl/profiling/question.py b/generalresearch/managers/thl/profiling/question.py index 7b2a7ad..10a9e32 100644 --- a/generalresearch/managers/thl/profiling/question.py +++ b/generalresearch/managers/thl/profiling/question.py @@ -1,6 +1,9 @@ +from __future__ import annotations + import random import threading -from typing import Any, Collection, Dict, List, Tuple +from collections.abc import Collection +from typing import Any from cachetools import TTLCache, cached from pydantic import ValidationError @@ -15,13 +18,13 @@ from generalresearch.models.thl.profiling.upk_question import ( class QuestionManager(PostgresManager): - def get_multi_upk(self, question_ids: Collection[str]) -> List[UpkQuestion]: + def get_multi_upk(self, question_ids: Collection[str]) -> list[UpkQuestion]: query = """ SELECT data, property_code, explanation_template, explanation_fragment_template FROM marketplace_question WHERE id = ANY(%(question_ids)s); """ - res: List[Dict[str, Any]] = self.pg_config.execute_sql_query( + res: list[dict[str, Any]] = self.pg_config.execute_sql_query( query=query, params={"question_ids": list(question_ids)} ) for x in res: @@ -41,7 +44,7 @@ class QuestionManager(PostgresManager): ) def get_questions_ranked( self, country_iso: str, language_iso: str - ) -> List[UpkQuestion]: + ) -> list[UpkQuestion]: query = """ SELECT data, property_code, explanation_template, explanation_fragment_template FROM marketplace_question @@ -51,11 +54,11 @@ class QuestionManager(PostgresManager): AND property_code NOT LIKE 'g:%%' AND is_live """ - res: List[Dict[str, Any]] = self.pg_config.execute_sql_query( + res: list[dict[str, Any]] = self.pg_config.execute_sql_query( query=query, params={"country_iso": country_iso, "language_iso": language_iso}, ) - qs: List[UpkQuestion] = [] + qs: list[UpkQuestion] = [] for x in res: x["data"]["ext_question_id"] = x["property_code"] x["data"]["explanation_template"] = x["explanation_template"] @@ -95,7 +98,7 @@ class QuestionManager(PostgresManager): "country_iso": country_iso, "language_iso": language_iso, } - res: List[Dict[str, Any]] = self.pg_config.execute_sql_query( + res: list[dict[str, Any]] = self.pg_config.execute_sql_query( query=query, params=params ) assert len(res) == 1, f"expected 1, got {len(res)} results" @@ -108,8 +111,8 @@ class QuestionManager(PostgresManager): return UpkQuestion.model_validate(x["data"]) def filter_by_property( - self, lookup: Collection[Tuple[str, str, str]] - ) -> List[UpkQuestion]: + self, lookup: Collection[tuple[str, str, str]] + ) -> list[UpkQuestion]: """ lookup is [(property_code, country_iso, language_iso)] """ @@ -123,7 +126,7 @@ class QuestionManager(PostgresManager): WHERE {where_str} """ flat_params = [item for tup in lookup for item in tup] - res: List[Dict[str, Any]] = self.pg_config.execute_sql_query( + res: list[dict[str, Any]] = self.pg_config.execute_sql_query( query=query, params=flat_params ) for x in res: @@ -141,7 +144,7 @@ class QuestionManager(PostgresManager): LOG.warning(e) return res2 - def update_question_explanation(self, q: UpkQuestion): + def update_question_explanation(self, q: UpkQuestion) -> None: # Assuming the question already exists in the db, and we're updating # the fields explanation_template and explanation_fragment_template assert q.id, "q.id must be set" @@ -160,4 +163,3 @@ class QuestionManager(PostgresManager): c.execute(query, params) assert c.rowcount == 1 conn.commit() - return None diff --git a/generalresearch/managers/thl/profiling/schema.py b/generalresearch/managers/thl/profiling/schema.py index e209067..1e3dece 100644 --- a/generalresearch/managers/thl/profiling/schema.py +++ b/generalresearch/managers/thl/profiling/schema.py @@ -1,5 +1,6 @@ +from __future__ import annotations + from threading import RLock -from typing import List from uuid import UUID from cachetools import TTLCache, cached @@ -13,7 +14,7 @@ from generalresearch.models.thl.profiling.upk_property import ( class UpkSchemaManager(PostgresManager): @cached(cache=TTLCache(maxsize=1, ttl=18 * 60), lock=RLock()) - def get_props_info(self) -> List[UpkProperty]: + def get_props_info(self) -> list[UpkProperty]: query = """ SELECT p.id AS property_id, @@ -68,7 +69,7 @@ class UpkSchemaManager(PostgresManager): c["id"] = UUID(c["id"]).hex return [UpkProperty.model_validate(x) for x in res] - def get_props_info_for_country(self, country_iso: str) -> List[UpkProperty]: + def get_props_info_for_country(self, country_iso: str) -> list[UpkProperty]: assert country_iso.lower() == country_iso res = self.get_props_info() res = [x for x in res if x.country_iso == country_iso].copy() diff --git a/generalresearch/managers/thl/profiling/uqa.py b/generalresearch/managers/thl/profiling/uqa.py index 1cab6c2..cbe39e7 100644 --- a/generalresearch/managers/thl/profiling/uqa.py +++ b/generalresearch/managers/thl/profiling/uqa.py @@ -1,6 +1,8 @@ +from __future__ import annotations + import logging +from collections.abc import Collection from datetime import datetime, timedelta, timezone -from typing import Collection, List, Optional from generalresearch.managers.base import PostgresManagerWithRedis from generalresearch.models.thl.profiling.user_question_answer import ( @@ -25,7 +27,7 @@ class UQAManager(PostgresManagerWithRedis): def update_cache( self, user: User, - uqas: List[UserQuestionAnswer], + uqas: list[UserQuestionAnswer], ): """ Adds new answers to the redis cache for this user. If the cache @@ -82,8 +84,8 @@ class UQAManager(PostgresManagerWithRedis): def _dedupe_and_clean_uqas( self, - uqas: List[UserQuestionAnswer], - ) -> List[UserQuestionAnswer]: + uqas: list[UserQuestionAnswer], + ) -> list[UserQuestionAnswer]: # Remove anything older than 30 days uqas = [uqa for uqa in uqas if not uqa.is_stale()] @@ -98,14 +100,14 @@ class UQAManager(PostgresManagerWithRedis): return sorted(new_uqas, key=lambda x: x.timestamp, reverse=True) - def get(self, user: User) -> List[UserQuestionAnswer]: + def get(self, user: User) -> list[UserQuestionAnswer]: uqas = self.get_from_cache(user=user) if uqas is None: uqas = self.recreate_cache(user) return self._dedupe_and_clean_uqas(uqas) - def get_from_cache(self, user: User) -> Optional[List[UserQuestionAnswer]]: + def get_from_cache(self, user: User) -> list[UserQuestionAnswer] | None: redis_key = self.redis_key(user) # Do the exists check and the list retrieval in a single transaction @@ -123,7 +125,7 @@ class UQAManager(PostgresManagerWithRedis): logger.info(f"{redis_key} exists") return uqas - def get_from_db(self, user: User) -> List[UserQuestionAnswer]: + def get_from_db(self, user: User) -> list[UserQuestionAnswer]: logger.info(f"get_uqa_from_db: {user.user_id}") # Only store the latest row per question_id. We don't need it multiple times. since = datetime.now(tz=timezone.utc) - timedelta(days=30) @@ -174,14 +176,12 @@ class UQAManager(PostgresManagerWithRedis): redis_key = f"thl-grpc:user-demographics:{user.user_id}" self.redis_client.delete(redis_key) - return None - def create( self, user: User, - uqas: List[UserQuestionAnswer], - session_id: Optional[str] = None, - ): + uqas: list[UserQuestionAnswer], + session_id: str | None = None, + ) -> None: for uqa in uqas: if uqa.user_id is None: uqa.user_id = user.user_id @@ -189,11 +189,10 @@ class UQAManager(PostgresManagerWithRedis): assert uqa.user_id == user.user_id self.create_in_db(uqas=uqas, session_id=session_id) self.update_cache(user=user, uqas=uqas) - return None def create_in_db( - self, uqas: List[UserQuestionAnswer], session_id: Optional[str] = None - ): + self, uqas: list[UserQuestionAnswer], session_id: str | None = None + ) -> None: values = [uqa.model_dump_mysql(session_id=session_id) for uqa in uqas] query = """ INSERT INTO marketplace_userquestionanswer @@ -208,4 +207,3 @@ class UQAManager(PostgresManagerWithRedis): with conn.cursor() as c: c.executemany(query=query, params_seq=values) conn.commit() - return None diff --git a/generalresearch/managers/thl/profiling/user_upk.py b/generalresearch/managers/thl/profiling/user_upk.py index b3ea52e..1820103 100644 --- a/generalresearch/managers/thl/profiling/user_upk.py +++ b/generalresearch/managers/thl/profiling/user_upk.py @@ -1,7 +1,10 @@ +from __future__ import annotations + import json from collections import defaultdict +from collections.abc import Collection from datetime import datetime, timedelta, timezone -from typing import Any, Collection, Dict, List, Optional, Set, Tuple, Union +from typing import Any from uuid import UUID from psycopg import Cursor @@ -29,8 +32,8 @@ class UserUpkManager(PostgresManagerWithRedis): self, pg_config: PostgresConfig, redis_config: RedisConfig, - permissions: Collection[Permission] = None, - cache_prefix: Optional[str] = None, + permissions: Collection[Permission] | None = None, + cache_prefix: str | None = None, ): super().__init__( pg_config=pg_config, @@ -42,9 +45,8 @@ class UserUpkManager(PostgresManagerWithRedis): def clear_upk_cache(self, user_id: int) -> None: self.redis_client.delete(f"thl-grpc:user-upk:{user_id}") - return None - def get_user_upk(self, user_id: int) -> List[UpkQuestionAnswer]: + def get_user_upk(self, user_id: int) -> list[UpkQuestionAnswer]: res = self.redis_client.get(f"thl-grpc:user-upk:{user_id}") if res: return [UpkQuestionAnswer.model_validate(x) for x in json.loads(res)] @@ -53,7 +55,7 @@ class UserUpkManager(PostgresManagerWithRedis): self.redis_client.set(f"thl-grpc:user-upk:{user_id}", value, ex=60 * 60 * 24) return res - def get_user_upk_mysql(self, user_id: int) -> List[UpkQuestionAnswer]: + def get_user_upk_mysql(self, user_id: int) -> list[UpkQuestionAnswer]: since = datetime.now(tz=timezone.utc) - timedelta(days=89) query = """ @@ -119,11 +121,11 @@ class UserUpkManager(PostgresManagerWithRedis): def get_user_upk_simple( self, user_id: PositiveInt, country_iso: str = "us" - ) -> Dict[str, Union[Set[str], str, float]]: + ) -> dict[str, set[str] | str | float]: res = self.get_user_upk(user_id=user_id) res = [x for x in res if x.country_iso == country_iso] - d: Dict[str, Union[Set[str], str, float]] = defaultdict(set) + d: dict[str, set[str] | str | float] = defaultdict(set) for x in res: if x.cardinality == Cardinality.ZERO_OR_ONE: d[x.property_label] = x.value @@ -134,7 +136,7 @@ class UserUpkManager(PostgresManagerWithRedis): def get_age_gender( self, user_id: PositiveInt, country_iso: str = "us" - ) -> Tuple[Optional[int], Optional[str]]: + ) -> tuple[int | None, str | None]: # Returns an integer year for age, and {'male', 'female', 'other_gender'} d = self.get_user_upk_simple(user_id, country_iso) @@ -145,14 +147,14 @@ class UserUpkManager(PostgresManagerWithRedis): gender = d.get("gender") return age, gender - def get_upk_schema(self, country_iso: str) -> List[UpkProperty]: + def get_upk_schema(self, country_iso: str) -> list[UpkProperty]: return self.upk_schema_manager.get_props_info_for_country( country_iso=country_iso ) def populate_user_upk_from_dict( - self, upk_ans_dict: List[Dict[str, Any]] - ) -> List[UpkQuestionAnswer]: + self, upk_ans_dict: list[dict[str, Any]] + ) -> list[UpkQuestionAnswer]: country_isos = {x["country_iso"] for x in upk_ans_dict} assert len(country_isos) == 1 @@ -193,7 +195,7 @@ class UserUpkManager(PostgresManagerWithRedis): upk_ans = [UpkQuestionAnswer.model_validate(x) for x in upk_ans_dict] return upk_ans - def upsert_user_profile_knowledge(self, c: Cursor, row: UpkQuestionAnswer): + def upsert_user_profile_knowledge(self, c: Cursor, row: UpkQuestionAnswer) -> None: prop_type_table = { PropertyType.UPK_ITEM: "marketplace_userprofileknowledgeitem", PropertyType.UPK_NUMERICAL: "marketplace_userprofileknowledgenumerical", @@ -244,7 +246,7 @@ class UserUpkManager(PostgresManagerWithRedis): def upsert_user_profile_knowledge_multi_item( self, c: Cursor, row: UpkQuestionAnswer - ): + ) -> None: args = row.model_dump_mysql() c.execute( @@ -299,9 +301,7 @@ class UserUpkManager(PostgresManagerWithRedis): args, ) - return None - - def set_user_upk(self, upk_ans: List[UpkQuestionAnswer]): + def set_user_upk(self, upk_ans: list[UpkQuestionAnswer]) -> None: user_id = {x.user_id for x in upk_ans} assert len(user_id) == 1, "only run for 1 user at a time" user_id = list(user_id)[0] |