Merge pull request #94 from balki/dev-2.x

Add typehint for IDE autocomplete
This commit is contained in:
Alex Root Junior 2019-01-01 16:21:46 +02:00 committed by GitHub
commit 9a4ed6a267
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

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