mirror of
https://github.com/aiogram/aiogram.git
synced 2025-12-09 01:15:31 +00:00
Return DataMixin
This commit is contained in:
parent
a823e275a7
commit
7db1572fd3
2 changed files with 62 additions and 2 deletions
|
|
@ -8,13 +8,39 @@ from typing import (
|
||||||
TypeVar,
|
TypeVar,
|
||||||
cast,
|
cast,
|
||||||
overload,
|
overload,
|
||||||
|
Dict,
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = ("ContextInstanceMixin",)
|
__all__ = ("ContextInstanceMixin", "DataMixin")
|
||||||
|
|
||||||
from typing_extensions import Literal
|
from typing_extensions import Literal
|
||||||
|
|
||||||
|
|
||||||
|
class DataMixin:
|
||||||
|
@property
|
||||||
|
def data(self) -> Dict[str, Any]:
|
||||||
|
data: Optional[Dict[str, Any]] = getattr(self, "_data", None)
|
||||||
|
if data is None:
|
||||||
|
data = {}
|
||||||
|
setattr(self, "_data", data)
|
||||||
|
return data
|
||||||
|
|
||||||
|
def __getitem__(self, key: str) -> Any:
|
||||||
|
return self.data[key]
|
||||||
|
|
||||||
|
def __setitem__(self, key: str, value: Any) -> None:
|
||||||
|
self.data[key] = value
|
||||||
|
|
||||||
|
def __delitem__(self, key: str) -> None:
|
||||||
|
del self.data[key]
|
||||||
|
|
||||||
|
def __contains__(self, key: str) -> bool:
|
||||||
|
return key in self.data
|
||||||
|
|
||||||
|
def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
|
||||||
|
return self.data.get(key, default)
|
||||||
|
|
||||||
|
|
||||||
ContextInstance = TypeVar("ContextInstance")
|
ContextInstance = TypeVar("ContextInstance")
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,45 @@
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from aiogram.utils.mixins import ContextInstanceMixin
|
from aiogram.utils.mixins import (
|
||||||
|
ContextInstanceMixin,
|
||||||
|
DataMixin,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ContextObject(ContextInstanceMixin['ContextObject']):
|
class ContextObject(ContextInstanceMixin['ContextObject']):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class DataObject(DataMixin):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class TestDataMixin:
|
||||||
|
def test_store_value(self):
|
||||||
|
obj = DataObject()
|
||||||
|
obj["foo"] = 42
|
||||||
|
|
||||||
|
assert "foo" in obj
|
||||||
|
assert obj["foo"] == 42
|
||||||
|
assert len(obj.data) == 1
|
||||||
|
|
||||||
|
def test_remove_value(self):
|
||||||
|
obj = DataObject()
|
||||||
|
obj["foo"] = 42
|
||||||
|
del obj["foo"]
|
||||||
|
|
||||||
|
assert "key" not in obj
|
||||||
|
assert len(obj.data) == 0
|
||||||
|
|
||||||
|
def test_getter(self):
|
||||||
|
obj = DataObject()
|
||||||
|
obj["foo"] = 42
|
||||||
|
|
||||||
|
assert obj.get("foo") == 42
|
||||||
|
assert obj.get("bar") is None
|
||||||
|
assert obj.get("baz", "test") == "test"
|
||||||
|
|
||||||
|
|
||||||
class TestContextInstanceMixin:
|
class TestContextInstanceMixin:
|
||||||
def test_empty(self):
|
def test_empty(self):
|
||||||
obj = ContextObject()
|
obj = ContextObject()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue