python-telegram-bot/telegram/ext/_utils/trackingdict.py
Bibo-Joshi 42eaa67fd5 Switch to asyncio (#2731)
Co-authored-by: tsnoam <tsnoam@gmail.com>
Co-authored-by: Harshil <37377066+harshil21@users.noreply.github.com>
Co-authored-by: Poolitzer <25934244+Poolitzer@users.noreply.github.com>
2022-05-06 18:19:15 +02:00

131 lines
4.5 KiB
Python

#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2022
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/].
"""This module contains a mutable mapping that keeps track of the keys that where accessed.
.. versionadded:: 14.0
Warning:
Contents of this module are intended to be used internally by the library and *not* by the
user. Changes to this module are not considered breaking changes and may not be documented in
the changelog.
"""
from typing import (
TypeVar,
Set,
ClassVar,
Union,
Tuple,
List,
Mapping,
Generic,
)
from collections import UserDict
from telegram._utils.defaultvalue import DEFAULT_NONE, DefaultValue
_VT = TypeVar('_VT')
_KT = TypeVar('_KT')
_T = TypeVar('_T')
class TrackingDict(UserDict, Generic[_KT, _VT]):
"""Mutable mapping that keeps track of which keys where accessed with write access.
Read-access is not tracked.
Note:
* ``setdefault()`` and ``pop`` are considered writing only depending on whether the
key is present
* deleting values is considered writing
"""
DELETED: ClassVar = object()
"""Special marker indicating that an entry was deleted."""
__slots__ = ('_write_access_keys',)
def __init__(self) -> None:
super().__init__()
self._write_access_keys: Set[_KT] = set()
def __track_write(self, key: Union[_KT, Set[_KT]]) -> None:
if isinstance(key, set):
self._write_access_keys |= key
else:
self._write_access_keys.add(key)
def pop_accessed_keys(self) -> Set[_KT]:
"""Returns all keys that were write-accessed since the last time this method was called."""
out = self._write_access_keys
self._write_access_keys = set()
return out
def pop_accessed_write_items(self) -> List[Tuple[_KT, _VT]]:
"""
Returns all keys & corresponding values as set of tuples that were write-accessed since
the last time this method was called. If a key was deleted, the value will be
:attr:`DELETED`.
"""
keys = self.pop_accessed_keys()
return [(key, self[key] if key in self else self.DELETED) for key in keys]
def mark_as_accessed(self, key: _KT) -> None:
"""Use this method have the key returned again in the next call to
:meth:`pop_accessed_write_items` or :meth:`pop_accessed_keys`
"""
self._write_access_keys.add(key)
# Override methods to track access
def __setitem__(self, key: _KT, value: _VT) -> None:
self.__track_write(key)
super().__setitem__(key, value)
def __delitem__(self, key: _KT) -> None:
self.__track_write(key)
super().__delitem__(key)
def update_no_track(self, mapping: Mapping[_KT, _VT]) -> None:
"""Like ``update``, but doesn't count towards write access."""
for key, value in mapping.items():
self.data[key] = value
# Mypy seems a bit inconsistent about what it wants as types for `default` and return value
# so we just ignore a bit
def pop( # type: ignore[override]
self, key: _KT, default: _VT = DEFAULT_NONE # type: ignore[assignment]
) -> _VT:
if key in self:
self.__track_write(key)
if isinstance(default, DefaultValue):
return super().pop(key)
return super().pop(key, default=default)
def clear(self) -> None:
self.__track_write(set(super().keys()))
super().clear()
# Mypy seems a bit inconsistent about what it wants as types for `default` and return value
# so we just ignore a bit
def setdefault(self: 'TrackingDict[_KT, _T]', key: _KT, default: _T = None) -> _T:
if key in self:
return self[key]
self.__track_write(key)
self[key] = default # type: ignore[assignment]
return default # type: ignore[return-value]