promnesia 1.2.20240810__py3-none-any.whl → 1.4.20250909__py3-none-any.whl

promnesia/server.py

@@ -1,38 +1,48 @@
-#!/usr/bin/python3
 from __future__ import annotations
 
 import argparse
-from dataclasses import dataclass
-from datetime import timedelta
-from functools import lru_cache
 import importlib.metadata
 import json
 import logging
 import os
+from dataclasses import dataclass
+from datetime import timedelta
+from functools import lru_cache
 from pathlib import Path
-from typing import List, NamedTuple, Dict, Optional, Any, Tuple, Protocol
-
-
-import pytz
-from pytz import BaseTzInfo
+from typing import Any, NamedTuple, Protocol
+from zoneinfo import ZoneInfo
 
 import fastapi
-
-from sqlalchemy import literal, between, or_, and_, exc, select
-from sqlalchemy import Column, Table, func, types
-from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy import (
+    Column,
+    Table,
+    and_,
+    between,
+    exc,
+    func,
+    literal,
+    or_,
+    select,
+    types,
+)
 from sqlalchemy.sql import text
+from sqlalchemy.sql.elements import ColumnElement
 
-
-from .common import PathWithMtime, DbVisit, Url, setup_logger, default_output_dir, get_system_tz
 from .cannon import canonify
+from .common import (
+    DbVisit,
+    PathWithMtime,
+    default_output_dir,
+    get_system_tz,
+    setup_logger,
+)
 from .database.load import DbStuff, get_db_stuff, row_to_db_visit
 
-
-Json = Dict[str, Any]
+Json = dict[str, Any]
 
 app = fastapi.FastAPI()
 
+
 # meh. need this since I don't have hooks in hug to initialize logging properly..
 @lru_cache(1)
 def get_logger() -> logging.Logger:
@@ -51,26 +61,26 @@ def get_logger() -> logging.Logger:
 
 
 def get_version() -> str:
+    assert __package__ is not None  # make type checker happy
     return importlib.metadata.version(__package__)
 
 
 class ServerConfig(NamedTuple):
     db: Path
-    timezone: BaseTzInfo
+    timezone: ZoneInfo
 
     def as_str(self) -> str:
-        return json.dumps({
-            'timezone': self.timezone.zone,
-            'db'      : str(self.db),
-        })
+        return json.dumps(
+            {
+                'timezone': self.timezone.key,
+                'db': str(self.db),
+            }
+        )
 
     @classmethod
-    def from_str(cls, cfgs: str) -> 'ServerConfig':
+    def from_str(cls, cfgs: str) -> ServerConfig:
         d = json.loads(cfgs)
-        return cls(
-            db      =Path         (d['db']),
-            timezone=pytz.timezone(d['timezone'])
-        )
+        return cls(db=Path(d['db']), timezone=ZoneInfo(d['timezone']))
 
 
 class EnvConfig:
@@ -88,8 +98,10 @@ class EnvConfig:
     def set(cfg: ServerConfig) -> None:
         os.environ[EnvConfig.KEY] = cfg.as_str()
 
+
 # todo how to return exception in error?
 
+
 def as_json(v: DbVisit) -> Json:
     # yep, this is NOT %Y-%m-%d as is seems to be the only format with timezone that Date.parse in JS accepts. Just forget it.
     dts = v.dt.strftime('%d %b %Y %H:%M:%S %z')
@@ -104,14 +116,14 @@ def as_json(v: DbVisit) -> Json:
         'duration': v.duration,
         'locator': {
             'title': loc.title,
-            'href' : loc.href,
+            'href': loc.href,
         },
-        'original_url'  : v.orig_url,
+        'original_url': v.orig_url,
         'normalised_url': v.norm_url,
     }
 
 
-def get_db_path(check: bool=True) -> Path:
+def get_db_path(*, check: bool = True) -> Path:
     db = EnvConfig.get().db
     if check:
         assert db.exists(), db
@@ -125,7 +137,7 @@ def _get_stuff(db_path: PathWithMtime) -> DbStuff:
     return get_db_stuff(db_path=db_path.path)
 
 
-def get_stuff(db_path: Optional[Path]=None) -> DbStuff: # TODO better name
+def get_stuff(db_path: Path | None = None) -> DbStuff:  # TODO better name
     # ok, it will always load from the same db file; but intermediate would be kinda an optional dump.
     if db_path is None:
         db_path = get_db_path()
@@ -136,15 +148,15 @@ def db_stats(db_path: Path) -> Json:
     engine, table = get_stuff(db_path)
     query = select(func.count()).select_from(table)
     with engine.connect() as conn:
-        total = list(conn.execute(query))[0][0]
+        [(total,)] = conn.execute(query)
     return {
         'total_visits': total,
     }
 
 
 class Where(Protocol):
-    def __call__(self, table: Table, url: str) -> ColumnElement[bool]:
-        ...
+    def __call__(self, table: Table, url: str) -> ColumnElement[bool]: ...
+
 
 @dataclass
 class VisitsResponse:
@@ -172,22 +184,21 @@ def search_common(url: str, where: Where) -> VisitsResponse:
     with engine.connect() as conn:
         try:
             # TODO make more defensive here
-            visits: List[DbVisit] = [row_to_db_visit(row) for row in conn.execute(query)]
+            visits: list[DbVisit] = [row_to_db_visit(row) for row in conn.execute(query)]
         except exc.OperationalError as e:
             if getattr(e, 'msg', None) == 'no such table: visits':
-                logger.warn('you may have to run indexer first!')
-                #result['visits'] = [{an error with a msg}] # TODO
-                #return result
+                logger.warning('you may have to run indexer first!')
+                # result['visits'] = [{an error with a msg}] # TODO
+                # return result
             raise
 
     logger.debug('got %d visits from db', len(visits))
 
-    vlist: List[DbVisit] = []
+    vlist: list[DbVisit] = []
     for vis in visits:
         dt = vis.dt
-        if dt.tzinfo is None: # FIXME need this for /visits endpoint as well?
-            tz = config.timezone
-            dt = tz.localize(dt)
+        if dt.tzinfo is None:  # FIXME need this for /visits endpoint as well?
+            dt = dt.replace(tzinfo=config.timezone)
             vis = vis._replace(dt=dt)
         vlist.append(vis)
 
@@ -202,8 +213,8 @@ def search_common(url: str, where: Where) -> VisitsResponse:
 
 # TODO hmm, seems that the extension is using post for all requests??
 # perhasp should switch to get for most endpoint
-@app.get ('/status', response_model=Json)
-@app.post('/status', response_model=Json)
+@app.get ('/status', response_model=Json)  # fmt: skip
+@app.post('/status', response_model=Json)  # fmt: skip
 def status() -> Json:
     '''
     Ideally, status will always respond, regardless the internal state of the backend?
@@ -225,7 +236,7 @@ def status() -> Json:
         logger.exception(e)
         stats = {'ERROR': str(e)}
 
-    version: Optional[str]
+    version: str | None
     try:
         version = get_version()
     except Exception as e:
@@ -236,15 +247,16 @@ def status() -> Json:
         'version': version,
         'db'     : db_path,
         'stats'  : stats,
-    }
+    }  # fmt: skip
 
 
 @dataclass
 class VisitsRequest:
     url: str
 
-@app.get ('/visits', response_model=VisitsResponse)
-@app.post('/visits', response_model=VisitsResponse)
+
+@app.get ('/visits', response_model=VisitsResponse)  # fmt: skip
+@app.post('/visits', response_model=VisitsResponse)  # fmt: skip
 def visits(request: VisitsRequest) -> VisitsResponse:
     url = request.url
     get_logger().info('/visited %s', url)
@@ -255,7 +267,7 @@ def visits(request: VisitsRequest) -> VisitsResponse:
             # exact match
             table.c.norm_url == url,
             # + child visits, but only 'interesting' ones
-            and_(table.c.context != None, table.c.norm_url.startswith(url, autoescape=True))  # noqa: E711
+            and_(table.c.context != None, table.c.norm_url.startswith(url, autoescape=True)),  # noqa: E711
         ),
     )
 
@@ -264,11 +276,13 @@ def visits(request: VisitsRequest) -> VisitsResponse:
 class SearchRequest:
     url: str
 
-@app.get ('/search', response_model=VisitsResponse)
-@app.post('/search', response_model=VisitsResponse)
+
+@app.get ('/search', response_model=VisitsResponse)  # fmt: skip
+@app.post('/search', response_model=VisitsResponse)  # fmt: skip
 def search(request: SearchRequest) -> VisitsResponse:
     url = request.url
     get_logger().info('/search %s', url)
+    # fmt: off
     return search_common(
         url=url,
         where=lambda table, url: or_(
@@ -279,71 +293,79 @@ def search(request: SearchRequest) -> VisitsResponse:
             table.c.locator_title.contains(url, autoescape=True),
         ),
     )
+    # fmt: on
 
 
 @dataclass
 class SearchAroundRequest:
     timestamp: float
 
-@app.get ('/search_around', response_model=VisitsResponse)
-@app.post('/search_around', response_model=VisitsResponse)
+
+@app.get ('/search_around', response_model=VisitsResponse)  # fmt: skip
+@app.post('/search_around', response_model=VisitsResponse)  # fmt: skip
 def search_around(request: SearchAroundRequest) -> VisitsResponse:
     timestamp = request.timestamp
     get_logger().info('/search_around %s', timestamp)
-    utc_timestamp = timestamp # old 'timestamp' name is legacy
+    utc_timestamp = timestamp  # old 'timestamp' name is legacy
 
     # TODO meh. use count/pagination instead?
-    delta_back  = timedelta(hours=3  ).total_seconds()
+    delta_back = timedelta(hours=3).total_seconds()
     delta_front = timedelta(minutes=2).total_seconds()
     # TODO not sure about delta_front.. but it also serves as quick hack to accommodate for all the truncations etc
 
     return search_common(
-        url='http://dummy.org', # NOTE: not used in the where query (below).. perhaps need to get rid of this
-        where=lambda table, url: between(
+        url='http://dummy.org',  # NOTE: not used in the where query (below).. perhaps need to get rid of this
+        where=lambda table, url: between(  # noqa: ARG005
             func.strftime(
-                '%s', # NOTE: it's tz aware, e.g. would distinguish +05:00 vs -03:00
+                '%s',  # NOTE: it's tz aware, e.g. would distinguish +05:00 vs -03:00
                 # this is a bit fragile, relies on cachew internal timestamp format, e.g.
                 # 2020-11-10T06:13:03.196376+00:00 Europe/London
                 func.substr(
                     table.c.dt,
-                    1, # substr is 1-indexed
+                    1,  # substr is 1-indexed
                     # instr finds the first match, but if not found it defaults to 0.. which we hack by concatting with ' '
                     func.instr(func.cast(table.c.dt, types.Unicode).op('||')(' '), ' ') - 1,
                     # for fucks sake.. seems that cast is necessary otherwise it tries to treat ' ' as datetime???
-                )
-            ) - literal(utc_timestamp),
+                ),
+            )
+            - literal(utc_timestamp),
             literal(-delta_back),
             literal(delta_front),
         ),
     )
 
+
 # before 0.11.14 (including), extension didn't share the version
 # so if it's not shared, assume that version
 _NO_VERSION = (0, 11, 14)
 _LATEST = (9999, 9999, 9999)
 
-def as_version(version: str) -> Tuple[int, int, int]:
+
+def as_version(version: str) -> tuple[int, int, int]:
     if version == '':
         return _NO_VERSION
     try:
         [v1, v2, v3] = map(int, version.split('.'))
-        return (v1, v2, v3)
     except Exception as e:
         logger = get_logger()
         logger.error('error while parsing version %s', version)
         logger.exception(e)
         return _LATEST
+    else:
+        return (v1, v2, v3)
 
 
 @dataclass
 class VisitedRequest:
-    urls: List[str]
+    urls: list[str]
     client_version: str = ''
 
-VisitedResponse = List[Optional[Json]]
 
-@app.get ('/visited', response_model=VisitedResponse)
-@app.post('/visited', response_model=VisitedResponse)
+VisitedResponse = list[Json | None]
+
+
+@app.get ('/visited', response_model=VisitedResponse)  # fmt: skip
+@app.post('/visited', response_model=VisitedResponse)  # fmt: skip
 def visited(request: VisitedRequest) -> VisitedResponse:
     # TODO instead switch logging to fastapi
     urls = request.urls
@@ -352,10 +374,10 @@ def visited(request: VisitedRequest) -> VisitedResponse:
     logger = get_logger()
     logger.info('/visited %s %s', urls, client_version)
 
-    version = as_version(client_version)
+    _version = as_version(client_version)  # todo use it?
 
     nurls = [canonify(u) for u in urls]
-    snurls = list(sorted(set(nurls)))
+    snurls = sorted(set(nurls))
 
     if len(snurls) == 0:
         return []
@@ -365,10 +387,11 @@ def visited(request: VisitedRequest) -> VisitedResponse:
     # sqlalchemy doesn't seem to support SELECT FROM (VALUES (...)) in its api
     # also doesn't support array binding...
     # https://stackoverflow.com/questions/13190392/how-can-i-bind-a-list-to-a-parameter-in-a-custom-query-in-sqlalchemy
-    bstring = ','.join(f'(:b{i})'   for i, _ in enumerate(snurls))
-    bdict = {            f'b{i}': v for i, v in enumerate(snurls)}
+    bstring = ','.join(f'(:b{i})'   for i, _ in enumerate(snurls))  # fmt: skip
+    bdict = {            f'b{i}': v for i, v in enumerate(snurls)}  # fmt: skip
     # TODO hopefully, visits.* thing only returns one visit??
-    query = text(f"""
+    query = (
+        text(f"""
 WITH cte(queried) AS (SELECT * FROM (values {bstring}))
 SELECT queried, visits.*
     FROM cte JOIN visits
@@ -378,9 +401,12 @@ SELECT queried, visits.*
     but somehow DESC is the one that actually works..
 */
     ORDER BY visits.context IS NULL DESC
-    """).bindparams(**bdict).columns(
-        Column('match', types.Unicode),
-        *table.columns,
+    """)
+        .bindparams(**bdict)
+        .columns(
+            Column('match', types.Unicode),
+            *table.columns,
+        )
     )
     # TODO might be very beneficial for performance to have an intermediate table
     # SELECT visits.* FROM visits GROUP BY visits.norm_url ORDER BY visits.context IS NULL DESC
@@ -388,10 +414,10 @@ SELECT queried, visits.*
     # brings down large queries to 50ms...
     with engine.connect() as conn:
         res = list(conn.execute(query))
-        present: Dict[str, Any] = {row[0]: row_to_db_visit(row[1:]) for row in res}
+        present: dict[str, Any] = {row[0]: row_to_db_visit(row[1:]) for row in res}
     results = []
     for nu in nurls:
-        r = present.get(nu, None)
+        r = present.get(nu)
         results.append(None if r is None else as_json(r))
 
     # no need for it anymore, extension has been updated since
@@ -411,6 +437,7 @@ def _run(*, host: str, port: str, quiet: bool, config: ServerConfig) -> None:
     EnvConfig.set(config)
 
     import uvicorn
+
     uvicorn.run('promnesia.server:app', host=host, port=int(port), log_level='debug')
 
 
@@ -422,7 +449,7 @@ def run(args: argparse.Namespace) -> None:
         config=ServerConfig(
             db=args.db,
             timezone=args.timezone,
-        )
+        ),
     )
 
 
@@ -464,7 +491,7 @@ def setup_parser(p: argparse.ArgumentParser) -> None:
 
     p.add_argument(
         '--timezone',
-        type=pytz.timezone,
+        type=ZoneInfo,
         default=get_system_tz(),
         help='Fallback timezone, defaults to the system timezone if not specified',
     )