tensorneko-util 0.3.14__py3-none-any.whl → 0.3.15__py3-none-any.whl

tensorneko_util/__init__.py

@@ -7,6 +7,7 @@ from . import visualization
 from . import preprocess
 from . import io
 from . import debug
+from . import msg
 from .io import read, write
 
 __all__ = [
@@ -16,9 +17,15 @@ __all__ = [
     "preprocess",
     "io",
     "debug",
+    "msg",
     "read",
     "write",
 ]
 
+try:
+    from . import _rich
+except ImportError:
+    pass
+
 with open(os.path.join(dirname(__file__), "version.txt"), "r", encoding="UTF-8") as file:
     __version__ = file.read()

tensorneko_util/_rich.py

@@ -0,0 +1,2 @@
+from rich.traceback import install
+install(show_locals=False)

tensorneko_util/io/json/json_writer.py

@@ -16,7 +16,8 @@ class JsonWriter:
             import orjson
             with open(path, "w", encoding=encoding) as file:
                 if indent == 4:
-                    option = orjson.OPT_INDENT_4
+                    warnings.warn("orjson does not support indent 4, will use indent 2 instead.")
+                    option = orjson.OPT_INDENT_2
                 elif indent == 2:
                     option = orjson.OPT_INDENT_2
                 else:

tensorneko_util/io/video/video_writer.py

@@ -1,5 +1,6 @@
 import pathlib
 import warnings
+from subprocess import Popen
 from typing import Optional, Union
 
 import numpy as np
@@ -116,7 +117,7 @@ class VideoWriter:
             if not VisualLib.ffmpeg_available():
                 raise ValueError("FFMPEG is not installed.")
             import ffmpeg
-            pipe = ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{video.shape[2]}x{video.shape[1]}") \
+            pipe: Popen = ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{video.shape[2]}x{video.shape[1]}") \
                 .output(path, format=ext[1:], pix_fmt="rgb24", r=video_fps) \
                 .overwrite_output() \
                 .run_async(pipe_stdin=True)

tensorneko_util/msg/__init__.py

@@ -0,0 +1,3 @@
+from .gotify import push_gotify
+
+__all__ = ["push_gotify"]

tensorneko_util/msg/gotify.py

@@ -0,0 +1,48 @@
+import json
+import os
+import socket
+import urllib.request
+from typing import Optional
+from urllib.error import HTTPError
+
+
+def push_gotify(message: str, url: Optional[str] = None, token: Optional[str] = None, title: Optional[str] = None,
+    priority: int = 0
+):
+    """
+    Push a message to a Gotify server.
+
+    Args:
+        message (``str``): The message to be sent.
+        url (``str``, optional): The URL of the Gotify server. If not provided, it will be read from the environment variable GOTIFY_URL.
+        token (``str``, optional): The token for the Gotify server. If not provided, it will be read from the environment variable GOTIFY_TOKEN.
+        title (``str``, optional): The title of the message. Default is the hostname of the machine.
+        priority (``int``, optional): The priority of the message. Default is 0.
+
+    Examples::
+
+        push_gotify("This is a test message", "<URL>", "<APP_TOKEN>")
+        # then the message will be sent to the Gotify server.
+        # title = "<HOST_NAME>", message = "This is a test message", priority = 0
+
+    """
+    title = title or socket.gethostname()
+    url = url or os.environ.get("GOTIFY_URL")
+    if url is None:
+        raise ValueError("URL is not provided. Input url argument or set the environment variable GOTIFY_URL.")
+    token = token or os.environ.get("GOTIFY_TOKEN")
+    if token is None:
+        raise ValueError("Token is not provided. Input token argument or set the environment variable GOTIFY_TOKEN.")
+    url = f"{url.rstrip('/')}/message?token={token}"
+    data = json.dumps(
+        {"title": title, "message": message, "priority": priority}
+    ).encode("utf-8")
+    req = urllib.request.Request(url, data=data, method="POST", headers={
+        "Content-Type": "application/json",
+        "User-Agent": "TensorNeko"
+    })
+    try:
+        with urllib.request.urlopen(req) as res:
+            pass
+    except HTTPError as e:
+        print(e.read().decode("utf-8"))

tensorneko_util/util/eventbus/bus.py

@@ -6,7 +6,8 @@ from abc import ABC, abstractmethod
 from concurrent.futures import wait
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, Type, Callable, TYPE_CHECKING, Any, Coroutine, Union, Set, TypeVar
+from typing import Dict, Type, Callable, TYPE_CHECKING, Any, Coroutine, Union, Set, TypeVar, List
+from concurrent.futures import Future
 
 from ...backend.parallel import ExecutorPool, ParallelType
 
@@ -99,12 +100,7 @@ class EventBus:
         .. code-block:: python
 
             # useful decorators for default event bus
-            from tensorneko.util import (
-                subscribe, # run in the main thread
-                subscribe_thread, # run in a new thread
-                subscribe_async, # run async
-                subscribe_process # run in a new process
-            )
+            from tensorneko.util import subscribe
             # Event base type
             from tensorneko.util import Event
 
@@ -117,15 +113,15 @@ class EventBus:
             def log_information(event: LogEvent):
                 print(event.message)
 
-            @subscribe_thread
+            @subscribe.thread
             def log_information_thread(event: LogEvent):
                 print(event.message, "in another thread")
 
-            @subscribe_async
+            @subscribe.coro
             async def log_information_async(event: LogEvent):
                 print(event.message, "async")
 
-            @subscribe_process
+            @subscribe.process
             def log_information_process(event: LogEvent):
                 print(event.message, "in a new process")
 
@@ -171,6 +167,7 @@ class EventBus:
 
     def __init__(self):
         self.listeners: Dict[Type[E], Set[Listener]] = {}
+        self.futures: List[Future] = []
 
     def register(self,
         event_type: Type[E],
@@ -185,7 +182,7 @@ class EventBus:
             self.listeners[event_type] = set()
         self.listeners[event_type].add(Listener(func, func_name, emit_type))
 
-    def emit(self, event: E, blocking: bool = True) -> None:
+    def emit(self, event: E) -> None:
         event_type: Type[E] = type(event)
         if event_type in self.listeners:
             normal_listeners = set()
@@ -216,8 +213,14 @@ class EventBus:
             for listener in normal_listeners:
                 listener.func(event)
 
-            if blocking:
-                wait(futures)
+            self.futures.extend(futures)
+
+    def wait(self):
+        """
+        Wait for all the futures to be completed.
+        """
+        wait(self.futures)
+        self.futures.clear()
 
     @staticmethod
     async def emit_async(async_listeners: Set[Listener], event: E) -> None:

tensorneko_util/util/eventbus/event.py

@@ -1,4 +1,3 @@
-from contextlib import contextmanager
 from typing import TypeVar
 
 from .bus import EventBus
@@ -9,24 +8,10 @@ class EventMeta(type):
     def __call__(cls, *args, **kwargs):
         event = super().__call__(*args, **kwargs)
         event.bus = kwargs.get("bus", EventBus.default)
-        event.bus.emit(event, blocking=_blocking_flag)
+        event.bus.emit(event)
         return event
 
 
-_blocking_flag = True
-
-
-@contextmanager
-def no_blocking():
-    """
-    Context manager to disable blocking event handler.
-    """
-    global _blocking_flag
-    _blocking_flag = False
-    yield
-    _blocking_flag = True
-
-
 class Event(metaclass=EventMeta):
     bus: EventBus
 

tensorneko_util/version.txt

@@ -1 +1 @@
-0.3.14
+0.3.15