plexus-python-pulse-common 1.1.75__py3-none-any.whl

plexus/pulse/common/utils/dockerutils.py

@@ -0,0 +1,181 @@
+import contextlib
+import dataclasses
+import re
+from collections.abc import Generator, Iterator
+from typing import Any
+
+import docker
+import docker.errors
+import docker.models.containers
+import docker.models.images
+from plexus.common.utils import logger
+from plexus.common.utils.jsonutils import JsonObject
+from plexus.common.utils.strutils import parse_int_or, trim_to_empty
+
+__all__ = [
+    "ImageName",
+    "docker_create_client",
+    "docker_build_image",
+    "docker_get_image",
+    "docker_pull_image",
+    "docker_fetch_image",
+    "docker_run_detached",
+]
+
+docker_image_tag_regexp: re.Pattern[str] = re.compile(
+    r"^(?P<components>[a-z0-9]+((__?|-+)[a-z0-9]+)*(/[a-z0-9]+((__?|-+)[a-z0-9]+)*)*)(:(?P<tag>\w[\w.-]{0,127}))?$")
+
+docker_image_tag_with_registry_regexp = re.compile(
+    r"^((?P<host>[a-zA-Z0-9.-]+)(:(?P<port>\d+))?/)?(?P<components>[a-z0-9]+((__?|-+)[a-z0-9]+)*(/[a-z0-9]+((__?|-+)[a-z0-9]+)*)*)(:(?P<tag>\w[\w.-]{0,127}))?$")
+
+
+@dataclasses.dataclass
+class ImageName(object):
+    registry_host: str | None
+    registry_port: int | None
+    components: list[str]
+    tag: str | None
+
+    @property
+    def registry(self) -> str:
+        if self.registry_host is None and self.registry_port is None:
+            return ""
+        if self.registry_port is None:
+            return self.registry_host
+        return f"{self.registry_host}:{self.registry_port}"
+
+    @property
+    def repository(self) -> str:
+        return "/".join(self.components)
+
+    @staticmethod
+    def parse(s: str):
+        match = docker_image_tag_regexp.match(s)
+        if match:
+            return ImageName(registry_host=None,
+                             registry_port=None,
+                             components=trim_to_empty(match.group("components")).split("/"),
+                             tag=match.group("tag"))
+
+        match = docker_image_tag_with_registry_regexp.match(s)
+        if match:
+            return ImageName(registry_host=match.group("host"),
+                             registry_port=parse_int_or(match.group("port"), None),
+                             components=trim_to_empty(match.group("components")).split("/"),
+                             tag=match.group("tag"))
+
+        return None
+
+
+def docker_create_client(
+    registry: str,
+    username: str,
+    password: str,
+) -> contextlib.AbstractContextManager[docker.DockerClient]:
+    client = docker.DockerClient()
+    client.login(registry=registry, username=username, password=password, reauth=True)
+    return contextlib.closing(client)
+
+
+def docker_build_image(
+    client: docker.DockerClient,
+    tag: str,
+    path: str,
+    dockerfile: str,
+    build_args: dict[str, str],
+) -> tuple[docker.models.images.Image, Iterator[dict[str, str]]]:
+    return client.images.build(tag=tag,
+                               path=path,
+                               dockerfile=dockerfile,
+                               buildargs=build_args,
+                               rm=True,
+                               forcerm=True,
+                               nocache=True)
+
+
+def docker_get_image(
+    client: docker.DockerClient,
+    image: str,
+) -> docker.models.images.Image:
+    return client.images.get(image)
+
+
+def docker_pull_image(
+    client: docker.DockerClient,
+    image: str,
+    fallback_local: bool = False,
+) -> docker.models.images.Image:
+    try:
+        return client.images.pull(image)
+    except docker.errors.ImageNotFound:
+        if not fallback_local:
+            raise
+        logger.warning("Image '%s' is not found from remote repository, try local repository instead", image)
+    except docker.errors.APIError:
+        if not fallback_local:
+            raise
+        logger.warning("Docker server returns an error while pulling image '%s', try local repository instead", image)
+
+    return docker_get_image(client, image)
+
+
+def docker_fetch_image(
+    client: docker.DockerClient,
+    image: str,
+    force_pull: bool = False,
+) -> docker.models.images.Image:
+    if force_pull:
+        return docker_pull_image(client, image, fallback_local=True)
+    try:
+        return docker_get_image(client, image)
+    except Exception:
+        return docker_pull_image(client, image, fallback_local=False)
+
+
+def docker_run_detached(
+    client: docker.DockerClient,
+    image: str,
+    name: str,
+    command: str | list[str],
+    volumes: dict[str, dict[str, str]],
+    environment: dict[str, str],
+    extra_hosts: dict[str, str],
+    timeout: int,
+    **kwargs,
+) -> tuple[JsonObject, Any]:
+    @contextlib.contextmanager
+    def managed_docker_run(
+        client: docker.DockerClient,
+        **kwargs,
+    ) -> Generator[docker.models.containers.Container, None, None]:
+        container_model = None
+        try:
+            container_model = client.containers.run(**kwargs)
+            yield container_model
+        finally:
+            if container_model is not None:
+                if container_model.status != "exited":
+                    try:
+                        container_model.stop()
+                    except docker.errors.DockerException:
+                        pass
+                container_model.wait()
+                try:
+                    container_model.remove()
+                except docker.errors.DockerException:
+                    pass
+
+    run_args = kwargs.copy()
+    run_args.update(dict(image=image,
+                         name=name,
+                         command=command,
+                         volumes=volumes,
+                         environment=environment,
+                         extra_hosts=extra_hosts,
+                         detach=True))
+
+    with managed_docker_run(client, **run_args) as container_model:
+        result = container_model.wait(timeout=timeout)
+        logs = container_model.logs()
+
+        return result, logs

plexus/pulse/common/utils/gisutils.py

@@ -0,0 +1,406 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import Annotated, Self
+
+import numpy as np
+import pydantic as pdt
+import pyproj
+import pyquaternion as pyquat
+from plexus.common.utils.funcutils import singleton
+from plexus.common.utils.strutils import parse_params_string, repr_data, str_conv
+
+
+class Pose(object):
+    @classmethod
+    def from_numbers(
+        cls,
+        px: float,
+        py: float,
+        pz: float,
+        qx: float,
+        qy: float,
+        qz: float,
+        qw: float,
+        ts: float = 0,
+    ) -> Self:
+        """
+        Constructs a pose from numbers representing position and orientation
+        """
+        return Pose(ts, np.array([px, py, pz], dtype=np.float64), np.array([qw, qx, qy, qz], dtype=np.float64))
+
+    @classmethod
+    def add(cls, x: Self, d: Self) -> Self:
+        """
+        Performs pose SE3 addition, as x + d = y
+        """
+        xq = pyquat.Quaternion(x.q)
+        dq = pyquat.Quaternion(d.q)
+
+        yp = x.p + xq.rotate(d.p)
+        yq = xq * dq
+
+        return Pose(0, yp, yq.normalised.elements)
+
+    @classmethod
+    def sub(cls, y: Self, x: Self) -> Self:
+        """
+        Performs pose SE3 subtraction, as x + d = y => d = y - x
+        """
+        xq = pyquat.Quaternion(x.q)
+        yq = pyquat.Quaternion(y.q)
+
+        dp = xq.inverse.rotate(y.p - x.p)
+        dq = xq.inverse * yq
+
+        return Pose(0, dp, dq.normalised.elements)
+
+    @classmethod
+    def interpolate(cls, a: Self, b: Self, t: float) -> Self:
+        """
+        Interpolates between two given poses, as a * t + b * (1 - t)
+
+        :return: interpolated pose
+        """
+        ts = a.ts + (b.ts - a.ts) * t
+        p = a.p + (b.p - a.p) * t
+        q = pyquat.Quaternion.slerp(pyquat.Quaternion(a.q), pyquat.Quaternion(b.q), t)
+        return Pose(ts, p, q.normalised.elements)
+
+    def __init__(self, ts: float, p: np.ndarray, q: np.ndarray):
+        """
+        Represents a pose
+
+        :param ts: timestamp
+        :param p: position vector
+        :param q: orientation quaternion
+        """
+        self.ts = ts
+        self.p = p
+        self.q = q
+
+    def matrix(self) -> np.ndarray:
+        """
+        Returns the transformation matrix of this pose
+
+        :return: transformation matrix
+        """
+        r = pyquat.Quaternion(self.q).rotation_matrix
+        t = self.p[:, None]
+        return np.block([[r, t], [np.zeros((1, 3), dtype=np.float64), np.ones((1, 1), dtype=np.float64)]])
+
+    def translate(self, v: np.ndarray) -> np.ndarray:
+        """
+        Translate the given vector with the pose position
+
+        :param v: the vector to be translated
+
+        :return: translated vector
+        """
+        return self.p + v
+
+    def rotate(self, v: np.ndarray) -> np.ndarray:
+        """
+        Rotates the given vector with the pose orientation
+
+        :param v: the vector to be rotated
+
+        :return: rotated vector
+        """
+        return pyquat.Quaternion(self.q).rotate(v)
+
+    def __str__(self):
+        return repr_data(self)
+
+
+class Proj(ABC):
+    @classmethod
+    @abstractmethod
+    def typename(cls) -> str:
+        pass
+
+    def __init__(self, spec):
+        self.spec = spec
+
+    def __str__(self):
+        return self.to_string()
+
+    def method(self) -> str:
+        return self.typename()
+
+    @abstractmethod
+    def ellipsoid(self) -> str:
+        pass
+
+    @abstractmethod
+    def from_latlon(self, lat: float, lon: float) -> tuple[float, float]:
+        """
+        Performs forward projection from latitude/longitude to projection coordinate system
+
+        :param lat: latitude
+        :param lon: longitude
+
+        :return: tuple of x and y in projection coordinate system
+        """
+        pass
+
+    @abstractmethod
+    def to_latlon(self, x: float, y: float) -> tuple[float, float]:
+        """
+        Performs inverse projection from projection coordinate system to latitude/longitude
+
+        :param x: x in projection coordinate system
+        :param y: y in projection coordinate system
+
+        :return: tuple of latitude and longitude
+        """
+        pass
+
+    @abstractmethod
+    def to_string(self) -> str:
+        pass
+
+
+class EQDCProj(Proj):
+    """
+    Equidistant conic projection
+    Please refer to https://en.wikipedia.org/wiki/Equidistant_conic_projection
+    """
+
+    class Spec(pdt.BaseModel):
+        lat_1: Annotated[float, pdt.Strict(), pdt.AllowInfNan(False)]
+        lat_2: Annotated[float, pdt.Strict(), pdt.AllowInfNan(False)]
+        lon_0: Annotated[float, pdt.Strict(), pdt.AllowInfNan(False)] = 0.0
+
+    @classmethod
+    def typename(cls) -> str:
+        return "eqdc"
+
+    def __init__(self, **kwargs):
+        super(EQDCProj, self).__init__(EQDCProj.Spec(**kwargs))
+        self.proj = self.make_proj()
+
+    def ellipsoid(self) -> str:
+        return "WGS84"
+
+    def from_latlon(self, lat: float, lon: float) -> tuple[float, float]:
+        x, y = self.proj(lon, lat)
+        return x, y
+
+    def to_latlon(self, x: float, y: float) -> tuple[float, float]:
+        lon, lat = self.proj(x, y, inverse=True)
+        return lat, lon
+
+    def to_string(self) -> str:
+        return "{method}:lat_1={},lat_2={},lon_0={}".format(
+            self.spec.lat_1,
+            self.spec.lat_2,
+            self.spec.lon_0,
+            method=self.method(),
+        )
+
+    def make_proj(self):
+        return pyproj.Proj(
+            "+proj={method} +ellps={ellipsoid} +datum={ellipsoid} +units=m +no_defs +lat_1={} +lat_2={} +lon_0={}".format(
+                self.spec.lat_1,
+                self.spec.lat_2,
+                self.spec.lon_0,
+                method=self.method(),
+                ellipsoid=self.ellipsoid(),
+            )
+        )
+
+
+class UTMProj(Proj):
+    """
+    Universal Transverse Mercator coordinate system
+    Please refer to https://en.wikipedia.org/wiki/Universal_Transverse_Mercator_coordinate_system
+    """
+
+    class Spec(pdt.BaseModel):
+        zone: Annotated[int, pdt.Strict()]
+        south: Annotated[bool, pdt.Strict()] = False
+
+    @classmethod
+    def typename(cls) -> str:
+        return "utm"
+
+    def __init__(self, **kwargs):
+        super(UTMProj, self).__init__(UTMProj.Spec(**kwargs))
+        self.proj = self.make_proj()
+
+    def ellipsoid(self) -> str:
+        return "WGS84"
+
+    def from_latlon(self, lat: float, lon: float) -> tuple[float, float]:
+        x, y = self.proj(lon, lat)
+        return x, y
+
+    def to_latlon(self, x: float, y: float) -> tuple[float, float]:
+        lon, lat = self.proj(x, y, inverse=True)
+        return lat, lon
+
+    def to_string(self) -> str:
+        return "{method}:zone={},south={}".format(self.spec.zone, self.spec.south, method=self.method())
+
+    def make_proj(self):
+        return pyproj.Proj(
+            "+proj={method} +ellps={ellipsoid} +datum={ellipsoid} +units=m +no_defs +zone={} {}".format(
+                self.spec.zone,
+                "+south" if self.spec.south else "",
+                method=self.method(),
+                ellipsoid=self.ellipsoid(),
+            )
+        )
+
+
+class WebMercProj(Proj):
+    """
+    Web Mercator projection
+    Please refer to https://en.wikipedia.org/wiki/Web_Mercator_projection
+    """
+
+    class Spec(pdt.BaseModel):
+        lon_0: Annotated[float, pdt.Strict(), pdt.AllowInfNan(False)] = 0.0
+
+    @classmethod
+    def typename(cls) -> str:
+        return "webmerc"
+
+    def __init__(self, **kwargs):
+        super(WebMercProj, self).__init__(WebMercProj.Spec(**kwargs))
+        self.proj = self.make_proj()
+
+    def ellipsoid(self) -> str:
+        return "WGS84"
+
+    def from_latlon(self, lat: float, lon: float) -> tuple[float, float]:
+        x, y = self.proj(lon, lat)
+        return x, y
+
+    def to_latlon(self, x: float, y: float) -> tuple[float, float]:
+        lon, lat = self.proj(x, y, inverse=True)
+        return lat, lon
+
+    def to_string(self) -> str:
+        return "{method}:lon_0={}".format(self.spec.lon_0, method=self.method())
+
+    def make_proj(self):
+        return pyproj.Proj(
+            "+proj={method} +ellps={ellipsoid} +datum={ellipsoid} +units=m +no_defs +lon_0={}".format(
+                self.spec.lon_0,
+                method=self.method(),
+                ellipsoid=self.ellipsoid(),
+            ),
+        )
+
+
+@singleton
+def default_proj() -> Proj:
+    return make_proj("eqdc:lat_1=30.0,lat_2=50.0")
+
+
+def make_proj(spec_str: str) -> Proj | None:
+    """
+    Constructs cartographic projection from the given name and arguments
+
+    :param spec_str: projection spec string
+
+    :return: projection
+    """
+
+    if spec_str is None:
+        return None
+
+    match spec_str.split(":", maxsplit=1):
+        case [name]:
+            kwargs = {}
+        case [name, args]:
+            kwargs = parse_params_string(args, str_parser=str_conv)
+        case _:
+            raise ValueError("bad spec string")
+
+    if name.lower() == UTMProj.typename():
+        return UTMProj(**kwargs)
+    if name.lower() == EQDCProj.typename():
+        return EQDCProj(**kwargs)
+    if name.lower() == WebMercProj.typename():
+        return WebMercProj(**kwargs)
+
+    return None
+
+
+class Coord(object):
+    @classmethod
+    def from_pose(cls, pose: Pose, proj: Proj | None) -> Self:
+        """
+        Creates a coord from given pose
+
+        :param pose: pose based on which the coord is created
+        :param proj: cartographic projection
+
+        :return: generated coord
+        """
+        return cls.from_xy(pose.p[0], pose.p[1], pose.p[2], proj, pose.ts)
+
+    @classmethod
+    def from_latlon(cls, lat: float, lon: float, ele: float, proj: Proj | None, ts: float = 0) -> Self:
+        """
+        Creates a coord from latitude/longitude
+
+        :param lat: latitude
+        :param lon: longitude
+        :param ele: elevation
+        :param proj: cartographic projection
+        :param ts: timestamp
+
+        :return: generated coord
+        """
+        if proj is None:
+            return Coord(ts, lat, lon, ele, 0.0, 0.0, pdt.BaseModel())
+
+        x, y = proj.from_latlon(lat, lon)
+        return Coord(ts, lat, lon, ele, x, y, proj.spec)
+
+    @classmethod
+    def from_xy(cls, x: float, y: float, ele: float, proj: Proj | None, ts: float = 0) -> Self:
+        """
+        Creates a coord from coordinate in cartographic projection
+
+        :param x: x coordinate in cartographic projection
+        :param y: y coordinate in cartographic projection
+        :param ele: elevation
+        :param proj: cartographic projection
+        :param ts: timestamp
+
+        :return: a coord
+        """
+        if proj is None:
+            return Coord(ts, 0.0, 0.0, ele, x, y, pdt.BaseModel())
+
+        lat, lon = proj.to_latlon(x, y)
+        return Coord(ts, lat, lon, ele, x, y, proj.spec)
+
+    def __init__(self, ts: float, lat: float, lon: float, ele: float, x: float, y: float, spec: pdt.BaseModel):
+        """
+        Represents a coordinate
+
+        :param ts: timestamp
+        :param lat: latitude
+        :param lon: longitude
+        :param ele: elevation
+        :param x: x in projection coordinate system
+        :param y: y in projection coordinate system
+        :param spec: projection specification
+        """
+        self.ts = ts
+        self.lat = lat
+        self.lon = lon
+        self.x = x
+        self.y = y
+        self.spec = spec
+        self.ele = ele
+        self.xy = np.array([x, y], dtype=np.float64)
+        self.xyz = np.array([x, y, ele], dtype=np.float64)
+
+    def __str__(self):
+        return repr_data(self)

plexus/pulse/common/utils/jsonutils.py

@@ -0,0 +1,96 @@
+import datetime
+import os
+from collections.abc import Generator, Iterable
+from typing import Any
+
+import ujson as json
+from plexus.common.utils.dtutils import dt_format, dt_parse, extended_format
+from plexus.common.utils.iterutils import batched
+from plexus.common.utils.jsonutils import JsonType, JsonValueCompatible
+from plexus.common.utils.jsonutils import json_reformat
+
+from plexus.pulse.common.utils.pathutils import collect_volumed_filenames, populate_volumed_filenames
+
+__all__ = [
+    "json_datetime_decoder",
+    "json_datetime_encoder",
+    "json_loads",
+    "json_dumps",
+    "read_chunked_jsonl",
+    "write_chunked_jsonl",
+]
+
+
+def json_datetime_decoder(v: Any) -> datetime.datetime | None:
+    if v is None:
+        return None
+    if isinstance(v, str):
+        return json_datetime_decoder(dt_parse(v, extended_format(with_us=True, with_tz=True)))
+    if isinstance(v, datetime.datetime):
+        return v.replace(tzinfo=datetime.timezone.utc)
+    raise ValueError("unexpected type of value for datetime decoder")
+
+
+def json_datetime_encoder(v: Any) -> str | None:
+    if v is None:
+        return None
+    if isinstance(v, str):
+        return json_datetime_encoder(dt_parse(v, extended_format(with_us=True, with_tz=True)))
+    if isinstance(v, datetime.datetime):
+        return dt_format(v.replace(tzinfo=datetime.timezone.utc), extended_format(with_us=True, with_tz=True))
+    raise ValueError("unexpected type of value for datetime encoder")
+
+
+def json_deserializer(obj):
+    def value_formatter(value: JsonValueCompatible) -> JsonType:
+        if not isinstance(value, str):
+            return value
+        try:
+            return dt_parse(value, extended_format(with_us=True, with_tz=True))
+        except Exception:
+            return value
+
+    return json_reformat(obj, value_formatter=value_formatter)
+
+
+def json_serializer(obj):
+    def unregistered_formatter(unregistered: Any) -> JsonType:
+        if isinstance(unregistered, datetime.datetime):
+            return dt_format(unregistered, extended_format(with_us=True, with_tz=True))
+        return None
+
+    return json_reformat(obj, raise_if_unregistered=False, unregistered_formatter=unregistered_formatter)
+
+
+def json_loads(s: str) -> JsonType:
+    return json_deserializer(json.loads(s))
+
+
+def json_dumps(obj: JsonType) -> str:
+    return json.dumps(json_serializer(obj), ensure_ascii=False, escape_forward_slashes=False)
+
+
+def read_chunked_jsonl(template: str) -> Generator[tuple[JsonType, str], None, None]:
+    for path, _ in collect_volumed_filenames(template):
+        with open(path, mode="r", encoding="utf-8") as fh:
+            for line in fh:
+                yield json_loads(line), path
+
+
+def write_chunked_jsonl(records: Iterable[JsonType], template: str, chunk_size: int) -> list[tuple[str, int]]:
+    generator = populate_volumed_filenames(template)
+    entry = []
+    for batch_index, batch in enumerate(batched(records, chunk_size)):
+        path, _ = next(generator)
+        lines = 0
+        with open(path, mode="w") as fh:
+            for record in batch:
+                fh.write(json_dumps(record))
+                fh.write("\n")
+                lines += 1
+        entry.append((path, lines))
+    if len(entry) == 1:
+        path, lines = entry[0]
+        os.rename(path, template)
+        return [(template, lines)]
+    return entry