From b58537bba4e954d49726e85b386acaa2bf6341c5 Mon Sep 17 00:00:00 2001
From: Ivan Melnychuk <ivan.melnychuk@pg3.ch>
Date: Tue, 29 Dec 2020 07:18:35 +0100
Subject: [PATCH] typing suport for context functions

---
 diagrams/__init__.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index ec57a607..f4e6430a 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -15,25 +15,25 @@ __diagram = contextvars.ContextVar("diagrams")
 __cluster = contextvars.ContextVar("cluster")
 
 
-def getdiagram():
+def getdiagram() -> "Diagram":
     try:
         return __diagram.get()
     except LookupError:
         return None
 
 
-def setdiagram(diagram):
+def setdiagram(diagram: "Diagram"):
     __diagram.set(diagram)
 
 
-def getcluster():
+def getcluster() -> "Cluster":
     try:
         return __cluster.get()
     except LookupError:
         return None
 
 
-def setcluster(cluster):
+def setcluster(cluster: "Cluster"):
     __cluster.set(cluster)