Add __getnewargs__ to SpanContext class (open-telemetry#1380)

rickychu · web-flow · commit f1db112cb0b4 · 2020-11-15T21:14:31.000-08:00
diff --git a/opentelemetry-api/CHANGELOG.md b/opentelemetry-api/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Unreleased
 
 - Add optional parameter to `record_exception` method ([#1314](https://github.com/open-telemetry/opentelemetry-python/pull/1314))
+- Add pickle support to SpanContext class ([#1380](https://github.com/open-telemetry/opentelemetry-python/pull/1380))
 
 ## Version 0.15b0
 
diff --git a/opentelemetry-api/src/opentelemetry/trace/span.py b/opentelemetry-api/src/opentelemetry/trace/span.py
@@ -190,6 +190,17 @@ def __new__(
             (trace_id, span_id, is_remote, trace_flags, trace_state, is_valid),
         )
 
+    def __getnewargs__(
+        self,
+    ) -> typing.Tuple[int, int, bool, "TraceFlags", "TraceState"]:
+        return (
+            self.trace_id,
+            self.span_id,
+            self.is_remote,
+            self.trace_flags,
+            self.trace_state,
+        )
+
     @property
     def trace_id(self) -> int:
         return self[0]  # pylint: disable=unsubscriptable-object
diff --git a/opentelemetry-api/tests/trace/test_span_context.py b/opentelemetry-api/tests/trace/test_span_context.py
@@ -0,0 +1,36 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pickle
+import unittest
+
+from opentelemetry import trace
+
+
+class TestSpanContext(unittest.TestCase):
+    def test_span_context_pickle(self):
+        """
+        SpanContext needs to be pickleable to support multiprocessing
+        so span can start as parent from the new spawned process
+        """
+        sc = trace.SpanContext(
+            1,
+            2,
+            is_remote=False,
+            trace_flags=trace.DEFAULT_TRACE_OPTIONS,
+            trace_state=trace.DEFAULT_TRACE_STATE,
+        )
+        pickle_sc = pickle.loads(pickle.dumps(sc))
+        self.assertEqual(sc.trace_id, pickle_sc.trace_id)
+        self.assertEqual(sc.span_id, pickle_sc.span_id)
