omlish 0.0.0.dev1__py3-none-any.whl → 0.0.0.dev2__py3-none-any.whl

omlish/procfs.py

@@ -0,0 +1,336 @@
+"""
+TODO:
+ - dataclasses
+"""
+import argparse
+import logging
+import os
+import re
+import resource
+import struct
+import sys
+import typing as ta
+
+from . import iterators as it
+from . import json
+from . import lang
+from . import os as oos
+
+
+log = logging.getLogger(__name__)
+
+
+PidLike = int | str
+
+
+RLIMIT_RESOURCES = {
+    getattr(resource, k): k
+    for k in dir(resource)
+    if k.startswith('RLIMIT_')
+}
+
+
+def parse_size(s: str) -> int:
+    if ' ' not in s:
+        return int(s)
+    us = {'kB': 1024, 'mB': 1024 * 1024}
+    v, u = s.split()
+    return int(v) * us[u]
+
+
+class ProcStat(lang.Namespace):
+    PID = 0
+    COMM = 1
+    STATE = 2
+    PPID = 3
+    PGRP = 4
+    SESSION = 5
+    TTY_NR = 6
+    TPGID = 7
+    FLAGS = 8
+    MINFLT = 9
+    CMINFLT = 10
+    MAJFLT = 11
+    CMAJFLT = 12
+    UTIME = 13
+    STIME = 14
+    CUTIME = 15
+    CSTIME = 16
+    PRIORITY = 17
+    NICE = 18
+    NUM_THREADS = 19
+    ITREALVALUE = 20
+    STARTTIME = 21
+    VSIZE = 22
+    RSS = 23
+    RSSLIM = 24
+    STARTCODE = 25
+    ENDCODE = 26
+    STARTSTACK = 27
+    KSTKESP = 28
+    KSTKEIP = 29
+    SIGNAL = 30
+    BLOCKED = 31
+    SIGIGNORE = 32
+    SIGCATCH = 33
+    WCHAN = 34
+    NSWAP = 35
+    CNSWAP = 36
+    EXIT_SIGNAL = 37
+    PROCESSOR = 38
+    RT_PRIORITY = 39
+    POLICY = 40
+    DELAYACCT_BLKIO_TICKS = 41
+    GUEST_TIME = 42
+    CGUEST_TIME = 43
+    START_DATA = 44
+    END_DATA = 45
+    START_BRK = 46
+    ARG_START = 47
+    ARG_END = 48
+    ENV_START = 49
+    ENV_END = 50
+    EXIT_CODE = 51
+
+
+def _check_linux() -> None:
+    if sys.platform != 'linux':
+        raise OSError
+
+
+def get_process_stats(pid: PidLike = 'self') -> ta.List[str]:
+    """http://man7.org/linux/man-pages/man5/proc.5.html -> /proc/[pid]/stat"""
+
+    _check_linux()
+    with open('/proc/%s/stat' % (pid,)) as f:
+        buf = f.read()
+    l, _, r = buf.rpartition(')')
+    pid, _, comm = l.partition('(')
+    return [pid.strip(), comm] + r.strip().split(' ')
+
+
+def get_process_chain(pid: PidLike = 'self') -> ta.List[ta.Tuple[int, str]]:
+    _check_linux()
+    lst = []
+    while pid:
+        process_stats = get_process_stats(pid)
+        lst.append((int(process_stats[ProcStat.PID]), process_stats[ProcStat.COMM]))
+        pid = int(process_stats[ProcStat.PPID])
+    return lst
+
+
+def get_process_start_time(pid: PidLike = 'self') -> int:
+    """https://stackoverflow.com/questions/2598145/how-to-retrieve-the-process-start-time-or-uptime-in-python"""
+
+    _check_linux()
+    hz = os.sysconf(os.sysconf_names['SC_CLK_TCK'])
+    with open('/proc/stat') as f:
+        system_stats = f.readlines()
+    for line in system_stats:
+        if line.startswith('btime'):
+            boot_timestamp = int(line.split()[1])
+            break
+    else:
+        raise ValueError
+    process_stats = get_process_stats(pid)
+    age_from_boot_jiffies = int(process_stats[ProcStat.STARTTIME])
+    age_from_boot_timestamp = age_from_boot_jiffies // hz
+    return boot_timestamp + age_from_boot_timestamp
+
+
+def get_process_rss(pid: PidLike = 'self') -> int:
+    return int(get_process_stats(pid)[ProcStat.RSS])
+
+
+def set_process_oom_score_adj(score: str, pid: PidLike = 'self') -> None:
+    _check_linux()
+    with open('/proc/%s/oom_score_adj' % (pid,), 'w') as f:
+        f.write(str(score))
+
+
+MAP_LINE_RX = re.compile(
+    r'^'
+    r'(?P<address>[A-Fa-f0-9]+)-(?P<end_address>[A-Fa-f0-9]+)\s+'
+    r'(?P<permissions>\S+)\s+'
+    r'(?P<offset>[A-Fa-f0-9]+)\s+'
+    r'(?P<device>\S+)\s+'
+    r'(?P<inode>\d+)\s+'
+    r'(?P<path>.*)'
+    r'$'
+)
+
+
+def get_process_maps(pid: PidLike = 'self', sharing: bool = False) -> ta.Iterator[ta.Dict[str, ta.Any]]:
+    """http://man7.org/linux/man-pages/man5/proc.5.html -> /proc/[pid]/maps"""
+
+    _check_linux()
+    with open('/proc/%s/%s' % (pid, 'smaps' if sharing else 'maps'), 'r') as map_file:
+        while True:
+            line = map_file.readline()
+            if not line:
+                break
+            m = MAP_LINE_RX.match(line)
+            if not m:
+                raise ValueError(line)
+            address = int(m.group('address'), 16)
+            end_address = int(m.group('end_address'), 16)
+            d = {
+                'address': address,
+                'end_address': end_address,
+                'size': end_address - address,
+                'permissions': [x for x in m.group('permissions') if x != '-'],
+                'offset': int(m.group('offset'), 16),
+                'device': m.group('device'),
+                'inode': int(m.group('inode')),
+                'path': m.group('path'),
+            }
+            if sharing:
+                s: dict[str, ta.Any] = {}
+                while True:
+                    line = map_file.readline()
+                    k, v = line.split(':')
+                    if k.lower() == 'vmflags':
+                        break
+                    s[k.lower()] = parse_size(v.strip())
+                _, v = line.split(':')
+                s['vmflags'] = [p for p in [j.strip() for j in v.split(' ')] if p]
+                d['sharing'] = s
+            yield d
+
+
+PAGEMAP_KEYS = (
+    'address',
+    'pfn',
+    'swap_type',
+    'swap_offset',
+    'pte_soft_dirty',
+    'file_page_or_shared_anon',
+    'page_swapped',
+    'page_present',
+)
+
+
+def get_process_range_pagemaps(start: int, end: int, pid: PidLike = 'self') -> ta.Iterable[ta.Dict[str, int]]:
+    """https://www.kernel.org/doc/Documentation/vm/pagemap.txt"""
+
+    _check_linux()
+    offset = (start // oos.PAGE_SIZE) * 8
+    npages = ((end - start) // oos.PAGE_SIZE)
+    size = npages * 8
+    with open('/proc/%s/pagemap' % (pid,), 'rb') as pagemap_file:
+        pagemap_file.seek(offset)
+        pagemap_buf = pagemap_file.read(size)
+    if not pagemap_buf:
+        return
+    _struct_unpack = struct.unpack
+    for pagenum in range(npages):
+        [packed] = _struct_unpack('Q', pagemap_buf[pagenum * 8:(pagenum + 1) * 8])
+        yield {
+            'address': start + (pagenum * oos.PAGE_SIZE),
+            'pfn': (packed & ((1 << (54 + 1)) - 1)),
+            'swap_type': (packed & ((1 << (4 + 1)) - 1)),
+            'swap_offset': (packed & ((1 << (54 + 1)) - 1)) >> 5,
+            'pte_soft_dirty': ((packed >> 55) & 1) > 0,
+            'file_page_or_shared_anon': ((packed >> 61) & 1) > 0,
+            'page_swapped': ((packed >> 62) & 1) > 0,
+            'page_present': ((packed >> 63) & 1) > 0,
+        }
+
+
+def get_process_pagemaps(pid: PidLike = 'self') -> ta.Iterable[ta.Dict[str, int]]:
+    _check_linux()
+    for m in get_process_maps(pid):
+        for p in get_process_range_pagemaps(m['address'], m['end_address'], pid):
+            yield p
+
+
+def _dump_cmd(args):
+    total = 0
+    dirty_total = 0
+    for m in get_process_maps(args.pid, sharing=True):
+        total += m['sharing']['rss']
+        sys.stdout.write(json.dumps({'map': m}))
+        sys.stdout.write('\n')
+        for pm in get_process_range_pagemaps(m['address'], m['end_address'], args.pid):
+            if pm['pte_soft_dirty']:
+                dirty_total += oos.PAGE_SIZE
+            sys.stdout.write(json.dumps({'page': tuple(pm[k] for k in PAGEMAP_KEYS)}))
+            sys.stdout.write('\n')
+    dct = {
+        'total': total,
+        'dirty_total': dirty_total,
+    }
+    sys.stdout.write(json.dumps(dct))
+    sys.stdout.write('\n')
+
+
+def _cmp_cmd(args):
+    if len(args.pids) == 1:
+        [rpid] = args.pids
+        lpid = get_process_chain(rpid)[1][0]
+    elif len(args.pids) == 2:
+        lpid, rpid = args.pids
+    else:
+        raise TypeError('Invalid arguments')
+
+    def g(pid):
+        for m in get_process_maps(pid, sharing=True):
+            for pm in get_process_range_pagemaps(m['address'], m['end_address'], pid):
+                yield pm
+
+    lpms, rpms = [g(pid) for pid in (lpid, rpid)]
+
+    l_pages = 0
+    r_pages = 0
+    c_pages = 0
+    for _, ps in it.merge_on(lambda pm: pm['address'], lpms, rpms):
+        l, r = it.expand_indexed_pairs(ps, None, width=2)
+        if l is not None and r is None:
+            l_pages += 1
+        elif l is None and r is not None:
+            r_pages += 1
+        elif l['pfn'] != r['pfn']:
+            c_pages += 1
+        else:
+            continue
+        if not args.quiet:
+            sys.stdout.write(json.dumps([l, r]))
+            sys.stdout.write('\n')
+    l_pages += c_pages
+    r_pages += c_pages
+    dct = {
+        'l_pages': l_pages,
+        'l_bytes': l_pages * oos.PAGE_SIZE,
+        'r_pages': r_pages,
+        'r_bytes': r_pages * oos.PAGE_SIZE,
+        'c_pages': c_pages,
+        'c_bytes': c_pages * oos.PAGE_SIZE,
+    }
+    sys.stdout.write(json.dumps(dct))
+    sys.stdout.write('\n')
+
+
+def _main():
+    _check_linux()
+
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument('-q', '--quiet', action='store_true')
+    arg_subparsers = arg_parser.add_subparsers()
+
+    dump_arg_parser = arg_subparsers.add_parser('dump')
+    dump_arg_parser.add_argument('pid', type=int)
+    dump_arg_parser.set_defaults(func=_dump_cmd)
+
+    cmp_arg_parser = arg_subparsers.add_parser('cmp')
+    cmp_arg_parser.add_argument('pids', type=int, nargs='*')
+    cmp_arg_parser.set_defaults(func=_cmp_cmd)
+
+    args = arg_parser.parse_args()
+    if not hasattr(args, 'func'):
+        arg_parser.print_help()
+    else:
+        args.func(args)
+
+
+if __name__ == '__main__':
+    _main()

omlish/reflect.py

@@ -63,16 +63,6 @@ _KNOWN_SPECIAL_TYPE_VARS = tuple(
 ##
 
 
-try:
-    from types import get_original_bases  # type: ignore
-except ImportError:
-    def get_original_bases(cls, /):
-        try:
-            return cls.__dict__.get('__orig_bases__', cls.__bases__)
-        except AttributeError:
-            raise TypeError(f'Expected an instance of type, not {type(cls).__name__!r}') from None
-
-
 def get_params(obj: ta.Any) -> tuple[ta.TypeVar, ...]:
     if isinstance(obj, type):
         if issubclass(obj, ta.Generic):  # type: ignore
@@ -218,7 +208,7 @@ def get_underlying(nt: NewType) -> Type:
     return type_(nt.obj.__supertype__)  # noqa
 
 
-def get_concrete_type(ty: Type) -> ta.Optional[type]:
+def get_concrete_type(ty: Type) -> type | None:
     if isinstance(ty, type):
         return ty
     if isinstance(ty, Generic):
@@ -299,7 +289,7 @@ class GenericSubstitution:
         if (cty := get_concrete_type(ty)) is not None:
             rpl = get_type_var_replacements(ty)
             ret: list[Type] = []
-            for b in get_original_bases(cty):
+            for b in types.get_original_bases(cty):
                 bty = type_(b)
                 if isinstance(bty, Generic) and isinstance(b, type):
                     # FIXME: throws away relative types, but can't use original vars as they're class-contextual

omlish/replserver/console.py

@@ -43,8 +43,8 @@ class InteractiveSocketConsole:
     def __init__(
             self,
             conn: sock.socket,
-            locals: ta.Optional[dict[str, ta.Any]] = None,
-            filename: str = '<console>'
+            locals: dict[str, ta.Any] | None = None,
+            filename: str = '<console>',
     ) -> None:
         super().__init__()
 
@@ -73,7 +73,7 @@ class InteractiveSocketConsole:
 
     CPRT = 'Type "help", "copyright", "credits" or "license" for more information.'
 
-    def interact(self, banner: ta.Optional[str] = None, exitmsg: ta.Optional[str] = None) -> None:
+    def interact(self, banner: str | None = None, exitmsg: str | None = None) -> None:
         log.info(f'Console {id(self)} on thread {threading.current_thread().ident} interacting')
 
         try:
@@ -144,10 +144,10 @@ class InteractiveSocketConsole:
 
     def compile(
             self,
-            source: ta.Union[str, ast.AST],
+            source: str | ast.AST,
             filename: str = '<input>',
-            symbol: str = 'single'
-    ) -> ta.Optional[types.CodeType]:
+            symbol: str = 'single',
+    ) -> types.CodeType | None:
         if isinstance(source, ast.AST):
             return self._compiler.compiler(source, filename, symbol)  # type: ignore
         else:
@@ -155,7 +155,7 @@ class InteractiveSocketConsole:
 
     def run_source(
             self,
-            source: ta.Union[str, ast.AST],
+            source: str | ast.AST,
             filename: str = '<input>',
             symbol: str = 'single',
     ) -> bool:
@@ -195,7 +195,7 @@ class InteractiveSocketConsole:
                         expr.value,
                         lineno=expr.lineno,
                         col_offset=expr.col_offset,
-                    )
+                    ),
                 ],
             )
             ast.fix_missing_locations(source)
@@ -227,7 +227,7 @@ class InteractiveSocketConsole:
         finally:
             last_tb = ei = None  # type: ignore  # noqa
 
-    def show_syntax_error(self, filename: ta.Optional[str] = None) -> None:
+    def show_syntax_error(self, filename: str | None = None) -> None:
         type, value, tb = sys.exc_info()
         sys.last_type = type
         sys.last_value = value

omlish/replserver/server.py

@@ -38,7 +38,7 @@ class ReplServer:
     @dc.dataclass(frozen=True)
     class Config:
         path: str
-        file_mode: ta.Optional[int] = None
+        file_mode: int | None = None
         poll_interval: float = 0.5
         exit_timeout: float = 10.0
 
@@ -51,7 +51,7 @@ class ReplServer:
         check.not_empty(config.path)
         self._config = check.isinstance(config, ReplServer.Config)
 
-        self._socket: ta.Optional[sock.socket] = None
+        self._socket: sock.socket | None = None
         self._is_running = False
         self._consoles_by_threads: ta.MutableMapping[threading.Thread, InteractiveSocketConsole] = \
             weakref.WeakKeyDictionary()  # noqa
@@ -106,7 +106,7 @@ class ReplServer:
                             log.info(
                                 f'Starting console {id(console)} repl server connection '
                                 f'on file {self._config.path} '
-                                f'on thread {threading.current_thread().ident}'
+                                f'on thread {threading.current_thread().ident}',
                             )
                             self._consoles_by_threads[threading.current_thread()] = console
                             console.interact()
@@ -135,7 +135,7 @@ class ReplServer:
                 self._is_shutdown.set()
                 self._is_running = False
 
-    def shutdown(self, block: bool = False, timeout: ta.Optional[float] = None) -> None:
+    def shutdown(self, block: bool = False, timeout: float | None = None) -> None:
         self._should_shutdown = True
         if block:
             self._is_shutdown.wait(timeout=timeout)

omlish/sql/_abcs.py

@@ -0,0 +1,65 @@
+import typing as ta
+
+
+DBAPITypeCode: ta.TypeAlias = ta.Any | None
+
+DBAPIColumnDescription: ta.TypeAlias = tuple[
+    str,
+    DBAPITypeCode,
+    int | None,
+    int | None,
+    int | None,
+    int | None,
+    bool | None,
+]
+
+
+class DBAPIConnection(ta.Protocol):
+    def close(self) -> object: ...
+
+    def commit(self) -> object: ...
+
+    # optional:
+    # def rollback(self) -> ta.Any: ...
+
+    def cursor(self) -> 'DBAPICursor': ...
+
+
+class DBAPICursor(ta.Protocol):
+    @property
+    def description(self) -> ta.Sequence[DBAPIColumnDescription] | None: ...
+
+    @property
+    def rowcount(self) -> int: ...
+
+    # optional:
+    # def callproc(self, procname: str, parameters: Sequence[ta.Any] = ...) -> Sequence[ta.Any]: ...
+
+    def close(self) -> object: ...
+
+    def execute(
+        self,
+        operation: str,
+        parameters: ta.Sequence[ta.Any] | ta.Mapping[str, ta.Any] = ...,
+    ) -> object: ...
+
+    def executemany(
+        self,
+        operation: str,
+        seq_of_parameters: ta.Sequence[ta.Sequence[ta.Any]],
+    ) -> object: ...
+
+    def fetchone(self) -> ta.Sequence[ta.Any] | None: ...
+
+    def fetchmany(self, size: int = ...) -> ta.Sequence[ta.Sequence[ta.Any]]: ...
+
+    def fetchall(self) -> ta.Sequence[ta.Sequence[ta.Any]]: ...
+
+    # optional:
+    # def nextset(self) -> None | Literal[True]: ...
+
+    arraysize: int
+
+    def setinputsizes(self, sizes: ta.Sequence[DBAPITypeCode | int | None]) -> object: ...
+
+    def setoutputsize(self, size: int, column: int = ...) -> object: ...

omlish/sql/dbs.py

@@ -0,0 +1,90 @@
+import typing as ta
+import urllib.parse
+
+from .. import dataclasses as dc
+from .. import lang
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class DbType:
+    name: str
+    dialect_name: str
+
+    default_port: int | None = None
+
+
+class DbTypes(lang.Namespace):
+    MYSQL = DbType(
+        name='mysql',
+        dialect_name='mysql',
+        default_port=3306,
+    )
+
+    POSTGRES = DbType(
+        name='postgres',
+        dialect_name='postgresql',
+        default_port=5432,
+    )
+
+    SQLITE = DbType(
+        name='sqlite',
+        dialect_name='sqlite',
+    )
+
+
+##
+
+
+class DbLoc(lang.Abstract):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class UrlDbLoc(DbLoc, lang.Final):
+    url: str
+
+
+@dc.dataclass(frozen=True)
+class HostDbLoc(DbLoc, lang.Final):
+    host: str
+    port: int | None = None
+
+    username: str | None = None
+    password: str | None = dc.xfield(default=None, repr_fn=lambda pw: '...' if pw is not None else None)
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class DbSpec:
+    name: str
+    type: DbType
+    loc: DbLoc
+
+
+##
+
+
+def rebuild_url(url: str, fn: ta.Callable[[urllib.parse.ParseResult], urllib.parse.ParseResult]) -> str:
+    if '://' in url:
+        engine, _, url = url.partition('://')
+        url = 'sql://' + url
+    else:
+        engine = None
+    parsed = urllib.parse.urlparse(url)
+    parsed = fn(parsed)
+    if engine is not None and parsed.scheme == 'sql':
+        parsed = parsed._replace(scheme=engine)
+    return urllib.parse.urlunparse(parsed)  # noqa
+
+
+def set_url_engine(url: str, engine: str) -> str:
+    return rebuild_url(url, lambda parsed: parsed._replace(scheme=engine))  # noqa
+
+
+def set_url_database(url: str, database: str) -> str:
+    return rebuild_url(url, lambda parsed: parsed._replace(path='/' + database))  # noqa

omlish/stats.py

@@ -272,10 +272,10 @@ class SamplingHistogram:
 
         self._percentile_pos_list = [self._calc_percentile_pos(p, self._size) for p in self._percentiles]
 
-        self._ring: list[ta.Optional[SamplingHistogram.Entry]] = [None] * size
+        self._ring: list[SamplingHistogram.Entry | None] = [None] * size
         self._ring_pos = 0
 
-        self._sample: list[ta.Optional[SamplingHistogram.Entry]] = [None] * size
+        self._sample: list[SamplingHistogram.Entry | None] = [None] * size
         self._sample_pos_queue = list(reversed(range(size)))
 
     def add(self, value: float) -> None:
@@ -303,7 +303,7 @@ class SamplingHistogram:
     def _calc_percentile_pos(p: float, sz: int) -> int:
         return int(round((p * sz) - 1))
 
-    def _calc_percentiles(self, entries: list[ta.Optional[Entry]]) -> list[Percentile]:
+    def _calc_percentiles(self, entries: list[Entry | None]) -> list[Percentile]:
         entries = list(filter(None, entries))
         sz = len(entries)
         if not sz:

omlish/testing/pydevd.py

@@ -5,8 +5,6 @@ an already-debugging PyCharm instance to debug PySpark jobs.
 
 TODO:
  - https://www.jetbrains.com/help/pycharm/remote-debugging-with-product.html#
- - move to dev?
- - cython help? or in cython.py
 """
 import json
 import os
@@ -87,7 +85,7 @@ def patch_for_trio_asyncio() -> None:
 
 
 @lang.cached_function
-def _pydevd() -> ta.Optional[types.ModuleType]:
+def _pydevd() -> types.ModuleType | None:
     try:
         return __import__('pydevd')
     except ImportError:
@@ -98,7 +96,7 @@ def is_present() -> bool:
     return _pydevd() is not None
 
 
-def get_setup() -> ta.Optional[dict]:
+def get_setup() -> dict | None:
     if is_present():
         return _pydevd().SetupHolder.setup
     else:
@@ -154,8 +152,8 @@ def save_args() -> None:
 
 def maybe_reexec(
         *,
-        file: ta.Optional[str] = None,
-        module: ta.Optional[str] = None,
+        file: str | None = None,
+        module: str | None = None,
         silence: bool = False,
 ) -> None:
     if ARGS_ENV_VAR not in os.environ:

omlish/testing/pytest/inject/__init__.py

@@ -1 +1,8 @@
 from . import harness  # noqa
+
+from .harness import (  # noqa
+    PytestScope,
+    Scopes,
+    bind,
+    register,
+)