Access to POSIX shared memory from Python

kitty/shm.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+# License: GPLv3 Copyright: 2022, Kovid Goyal <kovid at kovidgoyal.net>
+
+# This is present in the python stdlib after version 3.7 but we need to support
+# 3.7 for another year, so sigh.
+
+import mmap
+import os
+import secrets
+from typing import Optional
+
+from kitty.fast_data_types import shm_open, shm_unlink
+
+
+def make_filename(safe_length: int = 14, prefix: str = '/ky-') -> str:
+    "Create a random filename for the shared memory object."
+    # number of random bytes to use for name
+    nbytes = (safe_length - len(prefix)) // 2
+    name = prefix + secrets.token_hex(nbytes)
+    return name
+
+
+class SharedMemory:
+
+    def __init__(self, name: Optional[str] = None, create: bool = False, size: int = 0, readonly: bool = False, mode: int = 0o600):
+        if not size >= 0:
+            raise ValueError("'size' must be a positive integer")
+        if create:
+            flags = os.O_CREAT | os.O_EXCL
+            if size <= 0:
+                raise ValueError("'size' must be > 0")
+        else:
+            flags = os.O_RDONLY if readonly else os.O_RDWR
+        if name is None and not flags & os.O_EXCL:
+            raise ValueError("'name' can only be None if create=True")
+
+        if name is None:
+            while True:
+                name = make_filename()
+                try:
+                    self._fd = shm_open(name, flags, mode)
+                except FileExistsError:
+                    continue
+                self._name = name
+                break
+        self._name = name
+        try:
+            if create and size:
+                os.ftruncate(self._fd, size)
+            stats = os.fstat(self._fd)
+            size = stats.st_size
+            self._mmap = mmap.mmap(self._fd, size)
+        except OSError:
+            self.unlink()
+            raise
+
+        self.size = size
+        self._buf: Optional[memoryview] = memoryview(self._mmap)
+
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except OSError:
+            pass
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def buf(self) -> memoryview:
+        ans = self._buf
+        if ans is None:
+            raise RuntimeError('Cannot access the buffer of a closed shared memory object')
+        return ans
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.name!r}, size={self.size})'
+
+    def close(self) -> None:
+        """Closes access to the shared memory from this instance but does
+        not destroy the shared memory block."""
+        if self._buf is not None:
+            self._buf.release()
+            self._buf = None
+        if self._mmap is not None:
+            self._mmap.close()
+        if self._fd >= 0:
+            os.close(self._fd)
+            self._fd = -1
+
+    def unlink(self) -> None:
+        """Requests that the underlying shared memory block be destroyed.
+
+        In order to ensure proper cleanup of resources, unlink should be
+        called once (and only once) across all processes which have access
+        to the shared memory block."""
+        if self._name:
+            shm_unlink(self._name)
+            self._name = ''