meerschaum 2.9.5__py3-none-any.whl → 3.0.0__py3-none-any.whl

meerschaum/utils/daemon/RotatingFile.py

@@ -14,12 +14,13 @@ import traceback
 import sys
 import atexit
 from datetime import datetime, timezone
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Callable
 from meerschaum.config import get_config
 from meerschaum.utils.warnings import warn
 from meerschaum.utils.daemon.FileDescriptorInterceptor import FileDescriptorInterceptor
 from meerschaum.utils.threading import Thread
 import meerschaum as mrsm
+import threading
 daemon = mrsm.attempt_import('daemon')
 
 class RotatingFile(io.IOBase):
@@ -38,6 +39,7 @@ class RotatingFile(io.IOBase):
         redirect_streams: bool = False,
         write_timestamps: bool = False,
         timestamp_format: Optional[str] = None,
+        write_callback: Optional[Callable[[str], None]] = None,
     ):
         """
         Create a file-like object which manages other files.
@@ -66,6 +68,9 @@ class RotatingFile(io.IOBase):
         timestamp_format: str, default None
             If `write_timestamps` is `True`, use this format for the timestamps.
             Defaults to `'%Y-%m-%d %H:%M'`.
+
+        write_callback: Optional[Callable[[str], None]], default None
+            If provided, execute this callback with the data to be written.
         """
         self.file_path = pathlib.Path(file_path)
         if num_files_to_keep is None:
@@ -74,17 +79,18 @@ class RotatingFile(io.IOBase):
             max_file_size = get_config('jobs', 'logs', 'max_file_size')
         if timestamp_format is None:
             timestamp_format = get_config('jobs', 'logs', 'timestamps', 'format')
-        if num_files_to_keep < 2:
-            raise ValueError("At least 2 files must be kept.")
-        if max_file_size < 1:
-            raise ValueError("Subfiles must contain at least one byte.")
+        if num_files_to_keep < 1:
+            raise ValueError("At least 1 file must be kept.")
+        if max_file_size < 100:
+            raise ValueError("Subfiles must contain at least 100 bytes.")
 
         self.num_files_to_keep = num_files_to_keep
         self.max_file_size = max_file_size
         self.redirect_streams = redirect_streams
         self.write_timestamps = write_timestamps
         self.timestamp_format = timestamp_format
-        self.subfile_regex_pattern = re.compile(r'(.*)\.log(?:\.\d+)?$')
+        self.write_callback = write_callback
+        self.subfile_regex_pattern = re.compile(r'(.*)\.log(?:\.\d+)?')
 
         ### When subfiles are opened, map from their index to the file objects.
         self.subfile_objects = {}
@@ -186,7 +192,7 @@ class RotatingFile(io.IOBase):
         """
         try:
             return int(subfile_name.replace(self.file_path.name + '.', ''))
-        except Exception as e:
+        except Exception:
             return -1
 
 
@@ -272,7 +278,7 @@ class RotatingFile(io.IOBase):
                 try:
                     daemon.daemon.redirect_stream(sys.stdout, self._current_file_obj)
                     daemon.daemon.redirect_stream(sys.stderr, self._current_file_obj)
-                except OSError as e:
+                except OSError:
                     warn(
                         f"Encountered an issue when redirecting streams:\n{traceback.format_exc()}"
                     )
@@ -287,30 +293,36 @@ class RotatingFile(io.IOBase):
             self.is_subfile_too_large(latest_subfile_index, potential_new_len)
         )
         if create_new_file:
-            old_subfile_index = latest_subfile_index
-            new_subfile_index = old_subfile_index + 1
-            new_file_path = self.get_subfile_path_from_index(new_subfile_index)
-            self._previous_file_obj = self._current_file_obj
-            self._current_file_obj = open(new_file_path, 'a+', encoding='utf-8')
-            self.subfile_objects[new_subfile_index] = self._current_file_obj
-            self.flush()
-
-            if self._previous_file_obj is not None:
-                if self.redirect_streams:
-                    self._redirected_subfile_objects[old_subfile_index] = self._previous_file_obj
-                    daemon.daemon.redirect_stream(self._previous_file_obj, self._current_file_obj)
-                    daemon.daemon.redirect_stream(sys.stdout, self._current_file_obj)
-                    daemon.daemon.redirect_stream(sys.stderr, self._current_file_obj)
-                self.close(unused_only=True)
+            self.increment_subfiles()
 
-            ### Sanity check in case writing somehow fails.
-            if self._previous_file_obj is self._current_file_obj:
-                self._previous_file_obj = None
+        return self._current_file_obj
 
-            self.delete(unused_only=True)
+    def increment_subfiles(self, increment_by: int = 1):
+        """
+        Create a new subfile and switch the file pointer over.
+        """
+        latest_subfile_index = self.get_latest_subfile_index()
+        old_subfile_index = latest_subfile_index
+        new_subfile_index = old_subfile_index + increment_by
+        new_file_path = self.get_subfile_path_from_index(new_subfile_index)
+        self._previous_file_obj = self._current_file_obj
+        self._current_file_obj = open(new_file_path, 'a+', encoding='utf-8')
+        self.subfile_objects[new_subfile_index] = self._current_file_obj
+        self.flush()
 
-        return self._current_file_obj
+        if self.redirect_streams:
+            if self._previous_file_obj is not None:
+                self._redirected_subfile_objects[old_subfile_index] = self._previous_file_obj
+                daemon.daemon.redirect_stream(self._previous_file_obj, self._current_file_obj)
+            daemon.daemon.redirect_stream(sys.stdout, self._current_file_obj)
+            daemon.daemon.redirect_stream(sys.stderr, self._current_file_obj)
+        self.close(unused_only=True)
+
+        ### Sanity check in case writing somehow fails.
+        if self._previous_file_obj is self._current_file_obj:
+            self._previous_file_obj = None
 
+        self.delete(unused_only=True)
 
     def close(self, unused_only: bool = False) -> None:
         """
@@ -330,7 +342,7 @@ class RotatingFile(io.IOBase):
             try:
                 if not subfile_object.closed:
                     subfile_object.close()
-            except Exception as e:
+            except Exception:
                 warn(f"Failed to close an open subfile:\n{traceback.format_exc()}")
 
             _ = self.subfile_objects.pop(subfile_index, None)
@@ -359,6 +371,12 @@ class RotatingFile(io.IOBase):
         As such, if data is larger than max_file_size, then the corresponding subfile
         may exceed this limit.
         """
+        try:
+            if callable(self.write_callback):
+                self.write_callback(data)
+        except Exception:
+            warn(f"Failed to execute write callback:\n{traceback.format_exc()}")
+
         try:
             self.file_path.parent.mkdir(exist_ok=True, parents=True)
             if isinstance(data, bytes):
@@ -379,7 +397,7 @@ class RotatingFile(io.IOBase):
             except BrokenPipeError:
                 warn("BrokenPipeError encountered. The daemon may have been terminated.")
                 return
-            except Exception as e:
+            except Exception:
                 warn(f"Failed to write to subfile:\n{traceback.format_exc()}")
             self.flush()
             self.delete(unused_only=True)
@@ -410,11 +428,10 @@ class RotatingFile(io.IOBase):
         )
         for subfile_path_to_delete in existing_subfile_paths[0:end_ix]:
             subfile_index = self.get_index_from_subfile_name(subfile_path_to_delete.name)
-            subfile_object = self.subfile_objects.get(subfile_index, None)
 
             try:
                 subfile_path_to_delete.unlink()
-            except Exception as e:
+            except Exception:
                 warn(
                     f"Unable to delete subfile '{subfile_path_to_delete}':\n"
                     + f"{traceback.format_exc()}"
@@ -586,20 +603,21 @@ class RotatingFile(io.IOBase):
             if not subfile_object.closed:
                 try:
                     subfile_object.flush()
-                except Exception as e:
+                except Exception:
                     warn(f"Failed to flush subfile {subfile_index}:\n{traceback.format_exc()}")
+
         if self.redirect_streams:
             try:
                 sys.stdout.flush()
             except BrokenPipeError:
                 pass
-            except Exception as e:
+            except Exception:
                 warn(f"Failed to flush STDOUT:\n{traceback.format_exc()}")
             try:
                 sys.stderr.flush()
             except BrokenPipeError:
                 pass
-            except Exception as e:
+            except Exception:
                 warn(f"Failed to flush STDERR:\n{traceback.format_exc()}")
 
 
@@ -665,10 +683,19 @@ class RotatingFile(io.IOBase):
         for thread in interceptor_threads[:end_ix]:
             try:
                 thread.join()
-            except Exception as e:
+            except Exception:
                 warn(f"Failed to join interceptor threads:\n{traceback.format_exc()}")
         del interceptor_threads[:end_ix]
 
+    def touch(self):
+        """
+        Touch the latest subfile.
+        """
+        subfile_path = self.get_latest_subfile_path()
+        subfile_path.touch()
+
+    def isatty(self) -> bool:
+        return True
 
     def __repr__(self) -> str:
         """

meerschaum/utils/daemon/StdinFile.py

@@ -13,6 +13,7 @@ import os
 import selectors
 import traceback
 
+import meerschaum as mrsm
 from meerschaum.utils.typing import Optional, Union
 from meerschaum.utils.warnings import warn
 
@@ -25,6 +26,8 @@ class StdinFile(io.TextIOBase):
         self,
         file_path: Union[pathlib.Path, str],
         lock_file_path: Optional[pathlib.Path] = None,
+        decode: bool = True,
+        refresh_seconds: Union[int, float, None] = None,
     ):
         if isinstance(file_path, str):
             file_path = pathlib.Path(file_path)
@@ -38,6 +41,13 @@ class StdinFile(io.TextIOBase):
         self._file_handler = None
         self._fd = None
         self.sel = selectors.DefaultSelector()
+        self.decode = decode
+        self._write_fp = None
+        self._refresh_seconds = refresh_seconds
+
+    @property
+    def encoding(self):
+        return 'utf-8'
 
     @property
     def file_handler(self):
@@ -47,11 +57,9 @@ class StdinFile(io.TextIOBase):
         if self._file_handler is not None:
             return self._file_handler
 
-        if self.file_path.exists():
-            self.file_path.unlink()
-
-        self.file_path.parent.mkdir(parents=True, exist_ok=True)
-        os.mkfifo(self.file_path.as_posix(), mode=0o600)
+        if not self.file_path.exists():
+            self.file_path.parent.mkdir(parents=True, exist_ok=True)
+            os.mkfifo(self.file_path.as_posix(), mode=0o600)
 
         self._fd = os.open(self.file_path, os.O_RDONLY | os.O_NONBLOCK)
         self._file_handler = os.fdopen(self._fd, 'rb', buffering=0)
@@ -59,11 +67,19 @@ class StdinFile(io.TextIOBase):
         return self._file_handler
 
     def write(self, data):
+        if self._write_fp is None:
+            self.file_path.parent.mkdir(parents=True, exist_ok=True)
+            if not self.file_path.exists():
+                os.mkfifo(self.file_path.as_posix(), mode=0o600)
+            self._write_fp = open(self.file_path, 'wb')
+
         if isinstance(data, str):
             data = data.encode('utf-8')
-
-        with open(self.file_path, 'wb') as f:
-            f.write(data)
+        try:
+            self._write_fp.write(data)
+            self._write_fp.flush()
+        except BrokenPipeError:
+            pass
 
     def fileno(self):
         fileno = self.file_handler.fileno()
@@ -83,18 +99,19 @@ class StdinFile(io.TextIOBase):
                             self.blocking_file_path.unlink()
                     except Exception:
                         warn(traceback.format_exc())
-                    return data.decode('utf-8')
+                    return data.decode('utf-8') if self.decode else data
             except (OSError, EOFError):
                 pass
 
-            self.blocking_file_path.touch()
-            time.sleep(0.1)
+            if not self.blocking_file_path.exists():
+                self.blocking_file_path.touch()
+            time.sleep(self.refresh_seconds)
 
     def readline(self, size=-1):
-        line = ''
+        line = '' if self.decode else b''
         while True:
             data = self.read(1)
-            if not data or data == '\n':
+            if not data or ((data == '\n') if self.decode else (data == b'\n')):
                 break
             line += data
 
@@ -111,11 +128,34 @@ class StdinFile(io.TextIOBase):
             self._file_handler = None
             self._fd = None
 
+        if self._write_fp is not None:
+            try:
+                self._write_fp.close()
+            except BrokenPipeError:
+                pass
+            self._write_fp = None
+
+        try:
+            if self.blocking_file_path.exists():
+                self.blocking_file_path.unlink()
+        except Exception:
+            pass
         super().close()
 
     def is_open(self):
         return self._file_handler is not None
 
+    def isatty(self) -> bool:
+        return False
+
+    @property
+    def refresh_seconds(self) -> Union[int, float]:
+        """
+        How many seconds between checking for blocking functions.
+        """
+        if not self._refresh_seconds:
+            self._refresh_seconds = mrsm.get_config('system', 'cli', 'refresh_seconds')
+        return self._refresh_seconds
 
     def __str__(self) -> str:
         return f"StdinFile('{self.file_path}')"

meerschaum/utils/daemon/__init__.py

@@ -7,9 +7,16 @@ Manage Daemons via the `Daemon` class.
 """
 
 from __future__ import annotations
-import os, pathlib, shutil, json, datetime, threading, shlex
-from meerschaum.utils.typing import SuccessTuple, List, Optional, Callable, Any, Dict
-from meerschaum.config._paths import DAEMON_RESOURCES_PATH
+
+import os
+import pathlib
+import shutil
+import json
+import datetime
+import threading
+import shlex
+
+from meerschaum.utils.typing import SuccessTuple, List, Optional, Callable, Any, Dict, Union
 from meerschaum.utils.daemon.StdinFile import StdinFile
 from meerschaum.utils.daemon.Daemon import Daemon
 from meerschaum.utils.daemon.RotatingFile import RotatingFile
@@ -17,6 +24,27 @@ from meerschaum.utils.daemon.FileDescriptorInterceptor import FileDescriptorInte
 from meerschaum.utils.daemon._names import get_new_daemon_name
 
 
+__all__ = (
+    'daemon_action',
+    'daemon_entry',
+    'get_daemons',
+    'get_daemon_ids',
+    'get_running_daemons',
+    'get_stopped_daemons',
+    'get_paused_daemons',
+    'get_filtered_daemons',
+    'get_new_daemon_name',
+    'run_daemon',
+    'running_in_daemon',
+    'Daemon',
+    'StdinFile',
+    'RotatingFile',
+    'FileDescriptorInterceptor',
+)
+
+_daemons = {}
+
+
 def daemon_entry(sysargs: Optional[List[str]] = None) -> SuccessTuple:
     """Parse sysargs and execute a Meerschaum action as a daemon.
 
@@ -183,6 +211,7 @@ def get_daemon_ids() -> List[str]:
     """
     Return the IDs of all daemons on disk.
     """
+    from meerschaum.config._paths import DAEMON_RESOURCES_PATH
     return [
         daemon_dir
         for daemon_dir in sorted(os.listdir(DAEMON_RESOURCES_PATH))
@@ -270,8 +299,6 @@ def get_filtered_daemons(
             if warn:
                 _warn(f"Daemon '{d_id}' does not exist.", stack=False)
             continue
-        if d.hidden:
-            pass
         daemons.append(d)
     return daemons
 
@@ -280,6 +307,20 @@ def running_in_daemon() -> bool:
     """
     Return whether the current thread is running in a Daemon context.
     """
-    from meerschaum.config.static import STATIC_CONFIG
+    from meerschaum._internal.static import STATIC_CONFIG
     daemon_env_var = STATIC_CONFIG['environment']['daemon_id']
     return daemon_env_var in os.environ
+
+
+def get_current_daemon() -> Union[Daemon, None]:
+    """
+    If running withing a daemon context, return the corresponding `Daemon`.
+    Otherwise return `None`.
+    """
+    from meerschaum._internal.static import STATIC_CONFIG
+    daemon_env_var = STATIC_CONFIG['environment']['daemon_id']
+    daemon_id = os.environ.get(daemon_env_var, None)
+    if daemon_id is None:
+        return None
+
+    return _daemons.get(daemon_id, Daemon(daemon_id=daemon_id))

meerschaum/utils/daemon/_names.py

@@ -7,9 +7,11 @@ Generate random names for jobs.
 """
 
 from __future__ import annotations
-import os, random
-from meerschaum.utils.typing import Dict, List, Tuple
-from meerschaum.config._paths import DAEMON_RESOURCES_PATH
+
+import os
+import random
+
+from meerschaum.utils.typing import Dict, List
 
 _bank: Dict[str, Dict[str, List[str]]] = {
     'adjectives': {
@@ -116,6 +118,7 @@ def get_new_daemon_name() -> str:
     Generate a new random name until a unique one is found
     (up to ~6000 maximum possibilities).
     """
+    from meerschaum.config._paths import DAEMON_RESOURCES_PATH
     existing_names = (
         os.listdir(DAEMON_RESOURCES_PATH)
         if DAEMON_RESOURCES_PATH.exists()