Fix mypy issues.

open-telemetry · Aug 26, 2019 · 8e4d93f · 8e4d93f
1 parent 55a0fda
commit 8e4d93f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/opentelemetry-api/src/opentelemetry/distributedcontext/__init__.py b/opentelemetry-api/src/opentelemetry/distributedcontext/__init__.py
@@ -74,23 +74,26 @@ def __init__(
         self.value = value
 
 
-class DistributedContext(dict):
+class DistributedContext:
     """A container for distributed context entries"""
 
+    def __init__(self, entries: typing.Iterable[Entry]) -> None:
+        self._container = {entry.key: entry for entry in entries}
+
     def get_entries(self) -> typing.Iterable[Entry]:
         """Returns an immutable iterator to entries."""
-        return self.values()
+        return self._container.values()
 
     def get_entry_value(
             self,
             key: EntryKey
-    ) -> typing.Optional[EntryValue]:
+    ) -> typing.Optional[Entry]:
         """Returns the entry associated with a key or None
 
         Args:
             key: the key with which to perform a lookup
         """
-        return self.get(key)
+        return self._container.get(key)
 
 
 class DistributedContextManager:

diff --git a/opentelemetry-api/tests/distributedcontext/test_distributed_context.py b/opentelemetry-api/tests/distributedcontext/test_distributed_context.py
@@ -72,8 +72,9 @@ def setUp(self):
             distributedcontext.EntryKey("key"),
             distributedcontext.EntryValue("value"),
         )
-        context = self.context = distributedcontext.DistributedContext()
-        context[entry.key] = entry
+        context = self.context = distributedcontext.DistributedContext((
+            entry,
+        ))
 
     def test_get_entries(self):
         self.assertIn(self.entry, self.context.get_entries())

diff --git a/opentelemetry-sdk/tests/distributedcontext/test_distributed_context.py b/opentelemetry-sdk/tests/distributedcontext/test_distributed_context.py
@@ -27,7 +27,7 @@ def test_use_context(self):
         self.assertIsNone(self.manager.get_current_context())
 
         # Start initial context
-        dctx = dctx_api.DistributedContext()
+        dctx = dctx_api.DistributedContext(())
         with self.manager.use_context(dctx) as current:
             self.assertIs(current, dctx)
             self.assertIs(
@@ -36,7 +36,7 @@ def test_use_context(self):
             )
 
             # Context is overridden
-            nested_dctx = dctx_api.DistributedContext()
+            nested_dctx = dctx_api.DistributedContext(())
             with self.manager.use_context(nested_dctx) as current:
                 self.assertIs(current, nested_dctx)
                 self.assertIs(