streamlit-nightly 1.43.2.dev20250307__py3-none-any.whl

streamlit/runtime/media_file_manager.py

@@ -0,0 +1,234 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Provides global MediaFileManager object as `media_file_manager`."""
+
+from __future__ import annotations
+
+import collections
+import threading
+from typing import Final
+
+from streamlit.logger import get_logger
+from streamlit.runtime.media_file_storage import MediaFileKind, MediaFileStorage
+
+_LOGGER: Final = get_logger(__name__)
+
+
+def _get_session_id() -> str:
+    """Get the active AppSession's session_id."""
+    from streamlit.runtime.scriptrunner_utils.script_run_context import (
+        get_script_run_ctx,
+    )
+
+    ctx = get_script_run_ctx()
+    if ctx is None:
+        # This is only None when running "python myscript.py" rather than
+        # "streamlit run myscript.py". In which case the session ID doesn't
+        # matter and can just be a constant, as there's only ever "session".
+        return "dontcare"
+    else:
+        return ctx.session_id
+
+
+class MediaFileMetadata:
+    """Metadata that the MediaFileManager needs for each file it manages."""
+
+    def __init__(self, kind: MediaFileKind = MediaFileKind.MEDIA):
+        self._kind = kind
+        self._is_marked_for_delete = False
+
+    @property
+    def kind(self) -> MediaFileKind:
+        return self._kind
+
+    @property
+    def is_marked_for_delete(self) -> bool:
+        return self._is_marked_for_delete
+
+    def mark_for_delete(self) -> None:
+        self._is_marked_for_delete = True
+
+
+class MediaFileManager:
+    """In-memory file manager for MediaFile objects.
+
+    This keeps track of:
+    - Which files exist, and what their IDs are. This is important so we can
+      serve files by ID -- that's the whole point of this class!
+    - Which files are being used by which AppSession (by ID). This is
+      important so we can remove files from memory when no more sessions need
+      them.
+    - The exact location in the app where each file is being used (i.e. the
+      file's "coordinates"). This is is important so we can mark a file as "not
+      being used by a certain session" if it gets replaced by another file at
+      the same coordinates. For example, when doing an animation where the same
+      image is constantly replace with new frames. (This doesn't solve the case
+      where the file's coordinates keep changing for some reason, though! e.g.
+      if new elements keep being prepended to the app. Unlikely to happen, but
+      we should address it at some point.)
+    """
+
+    def __init__(self, storage: MediaFileStorage):
+        self._storage = storage
+
+        # Dict of [file_id -> MediaFileMetadata]
+        self._file_metadata: dict[str, MediaFileMetadata] = {}
+
+        # Dict[session ID][coordinates] -> file_id.
+        self._files_by_session_and_coord: dict[str, dict[str, str]] = (
+            collections.defaultdict(dict)
+        )
+
+        # MediaFileManager is used from multiple threads, so all operations
+        # need to be protected with a Lock. (This is not an RLock, which
+        # means taking it multiple times from the same thread will deadlock.)
+        self._lock = threading.Lock()
+
+    def _get_inactive_file_ids(self) -> set[str]:
+        """Compute the set of files that are stored in the manager, but are
+        not referenced by any active session. These are files that can be
+        safely deleted.
+
+        Thread safety: callers must hold `self._lock`.
+        """
+        # Get the set of all our file IDs.
+        file_ids = set(self._file_metadata.keys())
+
+        # Subtract all IDs that are in use by each session
+        for session_file_ids_by_coord in self._files_by_session_and_coord.values():
+            file_ids.difference_update(session_file_ids_by_coord.values())
+
+        return file_ids
+
+    def remove_orphaned_files(self) -> None:
+        """Remove all files that are no longer referenced by any active session.
+
+        Safe to call from any thread.
+        """
+        _LOGGER.debug("Removing orphaned files...")
+
+        with self._lock:
+            for file_id in self._get_inactive_file_ids():
+                file = self._file_metadata[file_id]
+                if file.kind == MediaFileKind.MEDIA:
+                    self._delete_file(file_id)
+                elif file.kind == MediaFileKind.DOWNLOADABLE:
+                    if file.is_marked_for_delete:
+                        self._delete_file(file_id)
+                    else:
+                        file.mark_for_delete()
+
+    def _delete_file(self, file_id: str) -> None:
+        """Delete the given file from storage, and remove its metadata from
+        self._files_by_id.
+
+        Thread safety: callers must hold `self._lock`.
+        """
+        _LOGGER.debug("Deleting File: %s", file_id)
+        self._storage.delete_file(file_id)
+        del self._file_metadata[file_id]
+
+    def clear_session_refs(self, session_id: str | None = None) -> None:
+        """Remove the given session's file references.
+
+        (This does not remove any files from the manager - you must call
+        `remove_orphaned_files` for that.)
+
+        Should be called whenever ScriptRunner completes and when a session ends.
+
+        Safe to call from any thread.
+        """
+        if session_id is None:
+            session_id = _get_session_id()
+
+        _LOGGER.debug("Disconnecting files for session with ID %s", session_id)
+
+        with self._lock:
+            if session_id in self._files_by_session_and_coord:
+                del self._files_by_session_and_coord[session_id]
+
+        _LOGGER.debug(
+            "Sessions still active: %r", self._files_by_session_and_coord.keys()
+        )
+
+        _LOGGER.debug(
+            "Files: %s; Sessions with files: %s",
+            len(self._file_metadata),
+            len(self._files_by_session_and_coord),
+        )
+
+    def add(
+        self,
+        path_or_data: bytes | str,
+        mimetype: str,
+        coordinates: str,
+        file_name: str | None = None,
+        is_for_static_download: bool = False,
+    ) -> str:
+        """Add a new MediaFile with the given parameters and return its URL.
+
+        If an identical file already exists, return the existing URL
+        and registers the current session as a user.
+
+        Safe to call from any thread.
+
+        Parameters
+        ----------
+        path_or_data : bytes or str
+            If bytes: the media file's raw data. If str: the name of a file
+            to load from disk.
+        mimetype : str
+            The mime type for the file. E.g. "audio/mpeg".
+            This string will be used in the "Content-Type" header when the file
+            is served over HTTP.
+        coordinates : str
+            Unique string identifying an element's location.
+            Prevents memory leak of "forgotten" file IDs when element media
+            is being replaced-in-place (e.g. an st.image stream).
+            coordinates should be of the form: "1.(3.-14).5"
+        file_name : str or None
+            Optional file_name. Used to set the filename in the response header.
+        is_for_static_download: bool
+            Indicate that data stored for downloading as a file,
+            not as a media for rendering at page. [default: False]
+
+        Returns
+        -------
+        str
+            The url that the frontend can use to fetch the media.
+
+        Raises
+        ------
+        If a filename is passed, any Exception raised when trying to read the
+        file will be re-raised.
+        """
+
+        session_id = _get_session_id()
+
+        with self._lock:
+            kind = (
+                MediaFileKind.DOWNLOADABLE
+                if is_for_static_download
+                else MediaFileKind.MEDIA
+            )
+            file_id = self._storage.load_and_get_id(
+                path_or_data, mimetype, kind, file_name
+            )
+            metadata = MediaFileMetadata(kind=kind)
+
+            self._file_metadata[file_id] = metadata
+            self._files_by_session_and_coord[session_id][coordinates] = file_id
+
+            return self._storage.get_url(file_id)

streamlit/runtime/media_file_storage.py

@@ -0,0 +1,143 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from enum import Enum
+from typing import Protocol
+
+
+class MediaFileKind(Enum):
+    # st.image, st.video, st.audio files
+    MEDIA = "media"
+
+    # st.download_button files
+    DOWNLOADABLE = "downloadable"
+
+
+class MediaFileStorageError(Exception):
+    """Exception class for errors raised by MediaFileStorage.
+
+    When running in "development mode", the full text of these errors
+    is displayed in the frontend, so errors should be human-readable
+    (and actionable).
+
+    When running in "release mode", errors are redacted on the
+    frontend; we instead show a generic "Something went wrong!" message.
+    """
+
+
+class MediaFileStorage(Protocol):
+    @abstractmethod
+    def load_and_get_id(
+        self,
+        path_or_data: str | bytes,
+        mimetype: str,
+        kind: MediaFileKind,
+        filename: str | None = None,
+    ) -> str:
+        """Load the given file path or bytes into the manager and return
+        an ID that uniquely identifies it.
+
+        It's an error to pass a URL to this function. (Media stored at
+        external URLs can be served directly to the Streamlit frontend;
+        there's no need to store this data in MediaFileStorage.)
+
+        Parameters
+        ----------
+        path_or_data
+            A path to a file, or the file's raw data as bytes.
+
+        mimetype
+            The media's mimetype. Used to set the Content-Type header when
+            serving the media over HTTP.
+
+        kind
+            The kind of file this is: either MEDIA, or DOWNLOADABLE.
+
+        filename : str or None
+            Optional filename. Used to set the filename in the response header.
+
+        Returns
+        -------
+        str
+            The unique ID of the media file.
+
+        Raises
+        ------
+        MediaFileStorageError
+            Raised if the media can't be loaded (for example, if a file
+            path is invalid).
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_url(self, file_id: str) -> str:
+        """Return a URL for a file in the manager.
+
+        Parameters
+        ----------
+        file_id
+            The file's ID, returned from load_media_and_get_id().
+
+        Returns
+        -------
+        str
+            A URL that the frontend can load the file from. Because this
+            URL may expire, it should not be cached!
+
+        Raises
+        ------
+        MediaFileStorageError
+            Raised if the manager doesn't contain an object with the given ID.
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_file(self, file_id: str) -> None:
+        """Delete a file from the manager.
+
+        This should be called when a given file is no longer referenced
+        by any connected client, so that the MediaFileStorage can free its
+        resources.
+
+        Calling `delete_file` on a file_id that doesn't exist is allowed,
+        and is a no-op. (This means that multiple `delete_file` calls with
+        the same file_id is not an error.)
+
+        Note: implementations can choose to ignore `delete_file` calls -
+        this function is a *suggestion*, not a *command*. Callers should
+        not rely on file deletion happening immediately (or at all).
+
+        Parameters
+        ----------
+        file_id
+            The file's ID, returned from load_media_and_get_id().
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        MediaFileStorageError
+            Raised if file deletion fails for any reason. Note that these
+            failures will generally not be shown on the frontend (file
+            deletion usually occurs on session disconnect).
+
+        """
+        raise NotImplementedError

streamlit/runtime/memory_media_file_storage.py

@@ -0,0 +1,181 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MediaFileStorage implementation that stores files in memory."""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import mimetypes
+import os.path
+from typing import Final, NamedTuple
+
+from streamlit.logger import get_logger
+from streamlit.runtime.media_file_storage import (
+    MediaFileKind,
+    MediaFileStorage,
+    MediaFileStorageError,
+)
+from streamlit.runtime.stats import CacheStat, CacheStatsProvider, group_stats
+
+_LOGGER: Final = get_logger(__name__)
+
+# Mimetype -> filename extension map for the `get_extension_for_mimetype`
+# function. We use Python's `mimetypes.guess_extension` for most mimetypes,
+# but (as of Python 3.9) `mimetypes.guess_extension("audio/wav")` returns None,
+# so we handle it ourselves.
+PREFERRED_MIMETYPE_EXTENSION_MAP: Final = {
+    "audio/wav": ".wav",
+    "text/vtt": ".vtt",
+}
+
+
+def _calculate_file_id(data: bytes, mimetype: str, filename: str | None = None) -> str:
+    """Hash data, mimetype, and an optional filename to generate a stable file ID.
+
+    Parameters
+    ----------
+    data
+        Content of in-memory file in bytes. Other types will throw TypeError.
+    mimetype
+        Any string. Will be converted to bytes and used to compute a hash.
+    filename
+        Any string. Will be converted to bytes and used to compute a hash.
+    """
+    filehash = hashlib.new("sha224", usedforsecurity=False)
+    filehash.update(data)
+    filehash.update(bytes(mimetype.encode()))
+
+    if filename is not None:
+        filehash.update(bytes(filename.encode()))
+
+    return filehash.hexdigest()
+
+
+def get_extension_for_mimetype(mimetype: str) -> str:
+    if mimetype in PREFERRED_MIMETYPE_EXTENSION_MAP:
+        return PREFERRED_MIMETYPE_EXTENSION_MAP[mimetype]
+
+    extension = mimetypes.guess_extension(mimetype, strict=False)
+    if extension is None:
+        return ""
+
+    return extension
+
+
+class MemoryFile(NamedTuple):
+    """A MediaFile stored in memory."""
+
+    content: bytes
+    mimetype: str
+    kind: MediaFileKind
+    filename: str | None
+
+    @property
+    def content_size(self) -> int:
+        return len(self.content)
+
+
+class MemoryMediaFileStorage(MediaFileStorage, CacheStatsProvider):
+    def __init__(self, media_endpoint: str):
+        """Create a new MemoryMediaFileStorage instance
+
+        Parameters
+        ----------
+        media_endpoint
+            The name of the local endpoint that media is served from.
+            This endpoint should start with a forward-slash (e.g. "/media").
+        """
+        self._files_by_id: dict[str, MemoryFile] = {}
+        self._media_endpoint = media_endpoint
+
+    def load_and_get_id(
+        self,
+        path_or_data: str | bytes,
+        mimetype: str,
+        kind: MediaFileKind,
+        filename: str | None = None,
+    ) -> str:
+        """Add a file to the manager and return its ID."""
+        file_data: bytes
+        if isinstance(path_or_data, str):
+            file_data = self._read_file(path_or_data)
+        else:
+            file_data = path_or_data
+
+        # Because our file_ids are stable, if we already have a file with the
+        # given ID, we don't need to create a new one.
+        file_id = _calculate_file_id(file_data, mimetype, filename)
+        if file_id not in self._files_by_id:
+            _LOGGER.debug("Adding media file %s", file_id)
+            media_file = MemoryFile(
+                content=file_data, mimetype=mimetype, kind=kind, filename=filename
+            )
+            self._files_by_id[file_id] = media_file
+
+        return file_id
+
+    def get_file(self, filename: str) -> MemoryFile:
+        """Return the MemoryFile with the given filename. Filenames are of the
+        form "file_id.extension". (Note that this is *not* the optional
+        user-specified filename for download files.)
+
+        Raises a MediaFileStorageError if no such file exists.
+        """
+        file_id = os.path.splitext(filename)[0]
+        try:
+            return self._files_by_id[file_id]
+        except KeyError as e:
+            raise MediaFileStorageError(
+                f"Bad filename '{filename}'. (No media file with id '{file_id}')"
+            ) from e
+
+    def get_url(self, file_id: str) -> str:
+        """Get a URL for a given media file. Raise a MediaFileStorageError if
+        no such file exists.
+        """
+        media_file = self.get_file(file_id)
+        extension = get_extension_for_mimetype(media_file.mimetype)
+        return f"{self._media_endpoint}/{file_id}{extension}"
+
+    def delete_file(self, file_id: str) -> None:
+        """Delete the file with the given ID."""
+        # We swallow KeyErrors here - it's not an error to delete a file
+        # that doesn't exist.
+        with contextlib.suppress(KeyError):
+            del self._files_by_id[file_id]
+
+    def _read_file(self, filename: str) -> bytes:
+        """Read a file into memory. Raise MediaFileStorageError if we can't."""
+        try:
+            with open(filename, "rb") as f:
+                return f.read()
+        except Exception as ex:
+            raise MediaFileStorageError(f"Error opening '{filename}'") from ex
+
+    def get_stats(self) -> list[CacheStat]:
+        # We operate on a copy of our dict, to avoid race conditions
+        # with other threads that may be manipulating the cache.
+        files_by_id = self._files_by_id.copy()
+
+        stats: list[CacheStat] = [
+            CacheStat(
+                category_name="st_memory_media_file_storage",
+                cache_name="",
+                byte_length=len(file.content),
+            )
+            for _, file in files_by_id.items()
+        ]
+        return group_stats(stats)

streamlit/runtime/memory_session_storage.py

@@ -0,0 +1,77 @@
+# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from cachetools import TTLCache
+
+from streamlit.runtime.session_manager import SessionInfo, SessionStorage
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+
+
+class MemorySessionStorage(SessionStorage):
+    """A SessionStorage that stores sessions in memory.
+
+    At most maxsize sessions are stored with a TTL of ttl seconds. This class is really
+    just a thin wrapper around cachetools.TTLCache that complies with the SessionStorage
+    protocol.
+    """
+
+    # NOTE: The defaults for maxsize and ttl are chosen arbitrarily for now. These
+    # numbers are reasonable as the main problems we're trying to solve at the moment are
+    # caused by transient disconnects that are usually just short network blips. In the
+    # future, we may want to increase both to support use cases such as saving state for
+    # much longer periods of time. For example, we may want session state to persist if
+    # a user closes their laptop lid and comes back to an app hours later.
+    def __init__(
+        self,
+        maxsize: int = 128,
+        ttl_seconds: int = 2 * 60,  # 2 minutes
+    ) -> None:
+        """Instantiate a new MemorySessionStorage.
+
+        Parameters
+        ----------
+        maxsize
+            The maximum number of sessions we allow to be stored in this
+            MemorySessionStorage. If an entry needs to be removed because we have
+            exceeded this number, either
+            - an expired entry is removed, or
+            - the least recently used entry is removed (if no entries have expired).
+
+        ttl_seconds
+            The time in seconds for an entry added to a MemorySessionStorage to live.
+            After this amount of time has passed for a given entry, it becomes
+            inaccessible and will be removed eventually.
+        """
+
+        self._cache: MutableMapping[str, SessionInfo] = TTLCache(
+            maxsize=maxsize, ttl=ttl_seconds
+        )
+
+    def get(self, session_id: str) -> SessionInfo | None:
+        return self._cache.get(session_id, None)
+
+    def save(self, session_info: SessionInfo) -> None:
+        self._cache[session_info.session.id] = session_info
+
+    def delete(self, session_id: str) -> None:
+        del self._cache[session_id]
+
+    def list(self) -> list[SessionInfo]:
+        return list(self._cache.values())