Extract profiling

src/dsdk/__init__.py

@@ -10,6 +10,7 @@
 from .postgres import Mixin as PostgresMixin
 from .postgres import Persistor as Postgres
 from .postgres import PredictionMixin as PostgresPredictionMixin
+from .profile import Profile, profile
 from .service import Batch, Delegate, Service, Task
 from .utils import (
     chunks,
@@ -19,7 +20,6 @@
     load_json_file,
     load_pickle_file,
     now_utc_datetime,
-    profile,
     retry,
 )
 
@@ -35,6 +35,7 @@
     "PostgresPredictionMixin",
     "PostgresMixin",
     "Postgres",
+    "Profile",
     "Service",
     "Task",
     "chunks",

src/dsdk/profile.py

@@ -1,11 +1,17 @@
 # -*- coding: utf-8 -*-
 """Profile."""
 
-from typing import Any, Dict, Optional
+from contextlib import contextmanager
+from logging import getLogger
+from time import perf_counter_ns
+from typing import Any, Dict, Generator, Optional
 
 from cfgenvy import yaml_type
 
 
+logger = getLogger(__name__)
+
+
 class Profile:
     """Profile."""
 
@@ -33,9 +39,40 @@ def _yaml_repr(cls, dumper, self, *, tag: str):
 
     def __init__(self, on: int, end: Optional[int] = None):
         """__init__."""
-        self.on = on
         self.end = end
+        self.on = on
 
     def as_yaml(self) -> Dict[str, Any]:
         """As yaml."""
         return {"end": self.end, "on": self.on}
+
+    def __repr__(self):
+        """__repr__."""
+        return f"Profile(end={self.end}, on={self.on})"
+
+    def __str__(self):
+        """__str__."""
+        return str(
+            {
+                "end": self.end,
+                "on": self.on,
+            }
+        )
+
+
+@contextmanager
+def profile(key: str) -> Generator[Profile, None, None]:
+    """Profile."""
+    # Replace return type with ContextManager[Profile] when mypy is fixed.
+    i = Profile(perf_counter_ns())
+    logger.info('{"key": "%s.on", "ns": "%s"}', key, i.on)
+    try:
+        yield i
+    finally:
+        i.end = perf_counter_ns()
+        logger.info(
+            '{"key": "%s.end", "ns": "%s", "elapsed": "%s"}',
+            key,
+            i.end,
+            i.end - i.on,
+        )

src/dsdk/utils.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-from contextlib import contextmanager
 from datetime import datetime, timezone, tzinfo
 from functools import wraps
 from json import dump as json_dump
@@ -12,14 +11,11 @@
 from pickle import dump as pickle_dump
 from pickle import load as pickle_load
 from sys import stderr, stdout
-from time import perf_counter_ns
 from time import sleep as default_sleep
-from typing import Any, Callable, Generator, Sequence
+from typing import Any, Callable, Sequence
 
 from dateutil import parser, tz
 
-from .profile import Profile
-
 
 logger = getLogger(__name__)
 
@@ -110,24 +106,6 @@ def now_utc_datetime() -> datetime:
     return datetime.now(tz=timezone.utc)
 
 
-@contextmanager
-def profile(key: str) -> Generator[Profile, None, None]:
-    """Profile."""
-    # Replace return type with ContextManager[Profile] when mypy is fixed.
-    i = Profile(perf_counter_ns())
-    logger.info('{"key": "%s.on", "ns": "%s"}', key, i.on)
-    try:
-        yield i
-    finally:
-        i.end = perf_counter_ns()
-        logger.info(
-            '{"key": "%s.end", "ns": "%s", "elapsed": "%s"}',
-            key,
-            i.end,
-            i.end - i.on,
-        )
-
-
 def retry(
     exceptions: Sequence[Exception],
     retries: int = 60,