From 5aa34625c6436a595d22ba933279a21eafc38eb5 Mon Sep 17 00:00:00 2001
From: balki <189196+balki@users.noreply.github.com>
Date: Mon, 31 Dec 2018 23:05:44 -0500
Subject: [PATCH] Add typehint for IDE autocomplete

---
 aiogram/utils/mixins.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/aiogram/utils/mixins.py b/aiogram/utils/mixins.py
index dba2f2cd..aa248cd9 100644
--- a/aiogram/utils/mixins.py
+++ b/aiogram/utils/mixins.py
@@ -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)