python/aiogram
1
0
Fork 0
mirror of https://github.com/aiogram/aiogram.git synced 2025-12-11 01:54:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add typehint for IDE autocomplete

Browse source
This commit is contained in:
balki 2018-12-31 23:05:44 -05:00
parent 9bc170231d
commit 5aa34625c6
1 changed files with 8 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

10
aiogram/utils/mixins.py
View file

@ -1,4 +1,7 @@
import contextvars import contextvars
from typing import TypeVar, Type
__all__ = ('DataMixin', 'ContextInstanceMixin')
class DataMixin: class DataMixin:
@ -23,19 +26,22 @@ class DataMixin:
return self.data.get(key, default) return self.data.get(key, default)
T = TypeVar('T')
class ContextInstanceMixin: class ContextInstanceMixin:
def __init_subclass__(cls, **kwargs): def __init_subclass__(cls, **kwargs):
cls.__context_instance = contextvars.ContextVar('instance_' + cls.__name__) cls.__context_instance = contextvars.ContextVar('instance_' + cls.__name__)
return cls return cls
@classmethod @classmethod
def get_current(cls, no_error=True): def get_current(cls: Type[T], no_error=True) -> T:
if no_error: if no_error:
return cls.__context_instance.get(None) return cls.__context_instance.get(None)
return cls.__context_instance.get() return cls.__context_instance.get()
@classmethod @classmethod
def set_current(cls, value): def set_current(cls: Type[T], value: T):
if not isinstance(value, cls): if not isinstance(value, cls):
raise TypeError(f"Value should be instance of '{cls.__name__}' not '{type(value).__name__}'") raise TypeError(f"Value should be instance of '{cls.__name__}' not '{type(value).__name__}'")
cls.__context_instance.set(value) cls.__context_instance.set(value)