From 4c634f348f8847fda139fe469e0e8adfabfd649a Mon Sep 17 00:00:00 2001 From: Chris Rossi Date: Wed, 12 Aug 2020 09:10:44 -0400 Subject: [PATCH] feat: use contextvars.ConvextVar instead of threading.local in Python 3 Closes #504 Co-authored-by: ndavydovdev --- google/cloud/ndb/context.py | 41 ++++++++++++++++++++++++++++++++----- setup.py | 4 ++++ 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/google/cloud/ndb/context.py b/google/cloud/ndb/context.py index 76ca071f..9784aaf3 100644 --- a/google/cloud/ndb/context.py +++ b/google/cloud/ndb/context.py @@ -26,12 +26,43 @@ from google.cloud.ndb import tasklets -class _LocalState(threading.local): - """Thread local state.""" +try: # pragma: NO PY2 COVER + import contextvars - def __init__(self): - self.context = None - self.toplevel_context = None + class _LocalState: + """Thread local state.""" + + def __init__(self): + self._toplevel_context = contextvars.ContextVar( + "_toplevel_context", default=None + ) + self._context = contextvars.ContextVar("_context", default=None) + + @property + def context(self): + return self._context.get() + + @context.setter + def context(self, value): + self._context.set(value) + + @property + def toplevel_context(self): + return self._toplevel_context.get() + + @toplevel_context.setter + def toplevel_context(self, value): + self._toplevel_context.set(value) + + +except ImportError: # pragma: NO PY3 COVER + + class _LocalState(threading.local): + """Thread local state.""" + + def __init__(self): + self.context = None + self.toplevel_context = None _state = _LocalState() diff --git a/setup.py b/setup.py index 9c1a5b4f..b8800d12 100644 --- a/setup.py +++ b/setup.py @@ -14,6 +14,7 @@ import io import os +import sys import setuptools @@ -28,6 +29,9 @@ def main(): "redis", ] + if sys.version_info.major == 3 and sys.version_info.minor < 7: + dependencies.append("contextvars") + setuptools.setup( name="google-cloud-ndb", version = "1.4.2",