Format stuff

AbhiPrasad · AbhiPrasad · commit e210aef748f0 · 2019-12-11T14:24:54.000-08:00
diff --git a/sentry_sdk/integrations/spark/spark_driver.py b/sentry_sdk/integrations/spark/spark_driver.py
@@ -13,9 +13,11 @@
 
     from sentry_sdk._types import Event, Hint
 
-class SparkTransactionContainer():
+
+class SparkTransactionContainer:
     sentry_spark_transaction = None
 
+
 class SparkIntegration(Integration):
     identifier = "spark"
 
@@ -46,21 +48,69 @@ def _sentry_patched_func(self, *args, **kwargs):
 
                 with Hub.current.start_span(op=op, description=name) as span:
                     if sparkContext:
-                        span.set_tag("callsite", sparkContext.getLocalProperty("callSite.short"))
+                        span.set_tag(
+                            "callsite", sparkContext.getLocalProperty("callSite.short")
+                        )
                     return func(self, *args, **kwargs)
-            
+
             return _sentry_patched_func
 
         setattr(obj, name, wrap_rdd_func_call(rdd_func))
 
+
 def patch_spark_for_spans():
     # type: () -> None
     from pyspark import SparkContext
     from pyspark import RDD
 
-    SPARK_FUNCS = frozenset(["emptyRDD", "parallelize", "wholeTextFiles", "binaryFiles", "binaryFiles", "sequenceFile", "newAPIHadoopFile", ])
-    RDD_ACTIONS = frozenset(["reduce", "collect", "first", "take", "takeSample", "takeOrdered", "saveAsTextFile", "saveAsSequenceFile", "saveAsObjectFile", "countByKey", "foreach"])
-    RDD_TRANSFORMATIONS = frozenset(["map", "filter", "flatMap", "sample", "union", "intersection", "distinct", "groupByKey", "reduceByKey", "aggregateByKey", "sortByKey", "join", "cogroup", "cartesian", "pipe", "coalesce", "repartition", "repartitionAndSortWithinPartitions"])
+    SPARK_FUNCS = frozenset(
+        [
+            "emptyRDD",
+            "parallelize",
+            "wholeTextFiles",
+            "binaryFiles",
+            "binaryFiles",
+            "sequenceFile",
+            "newAPIHadoopFile",
+        ]
+    )
+    RDD_ACTIONS = frozenset(
+        [
+            "reduce",
+            "collect",
+            "first",
+            "take",
+            "takeSample",
+            "takeOrdered",
+            "saveAsTextFile",
+            "saveAsSequenceFile",
+            "saveAsObjectFile",
+            "countByKey",
+            "foreach",
+        ]
+    )
+    RDD_TRANSFORMATIONS = frozenset(
+        [
+            "map",
+            "filter",
+            "flatMap",
+            "sample",
+            "union",
+            "intersection",
+            "distinct",
+            "groupByKey",
+            "reduceByKey",
+            "aggregateByKey",
+            "sortByKey",
+            "join",
+            "cogroup",
+            "cartesian",
+            "pipe",
+            "coalesce",
+            "repartition",
+            "repartitionAndSortWithinPartitions",
+        ]
+    )
 
     wrap_span_around_func(SparkContext, SPARK_FUNCS, "spark.createRDD")
     wrap_span_around_func(RDD, RDD_ACTIONS, "spark.action")
@@ -76,7 +126,9 @@ def patch_spark_context_stop(spark_transaction_container):
     def _sentry_patched_spark_context_stop(self, *args, **kwargs):
         if spark_transaction_container.span:
             spark_transaction_container.span._tags.setdefault("app_name", self.appName)
-            spark_transaction_container.span._tags.setdefault("application_id", self.applicationId)
+            spark_transaction_container.span._tags.setdefault(
+                "application_id", self.applicationId
+            )
 
         stop = spark_context_stop(self, *args, **kwargs)
 
@@ -87,6 +139,7 @@ def _sentry_patched_spark_context_stop(self, *args, **kwargs):
 
     SparkContext.stop = _sentry_patched_spark_context_stop
 
+
 def _set_app_properties():
     # type: () -> None
     """
@@ -129,7 +182,10 @@ def _sentry_patched_spark_context_init(self, *args, **kwargs):
         if Hub.current.get_integration(SparkIntegration) is None:
             return init
 
-        spark_transaction_container.span = Hub.current.start_span(op="spark.job", transaction="{} - {}".format(self.appName, self.applicationId)).__enter__()
+        spark_transaction_container.span = Hub.current.start_span(
+            op="spark.job",
+            transaction="{} - {}".format(self.appName, self.applicationId),
+        ).__enter__()
 
         _start_sentry_listener(self)
         _set_app_properties()
@@ -140,7 +196,10 @@ def _sentry_patched_spark_context_init(self, *args, **kwargs):
             def process_event(event, hint):
                 # type: (Event, Hint) -> Optional[Event]
                 with capture_internal_exceptions():
-                    if Hub.current.get_integration(SparkIntegration) is None or not self._jsc:
+                    if (
+                        Hub.current.get_integration(SparkIntegration) is None
+                        or not self._jsc
+                    ):
                         return event
 
                     event.setdefault("user", {}).setdefault("id", self.sparkUser())
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -314,7 +314,7 @@ def finish(self, hub=None):
 
         if self._span_recorder is None:
             return None
-        
+
         print("LETS _span_recorder")
         print(self)
 
@@ -324,7 +324,7 @@ def finish(self, hub=None):
             # If this has no transaction set we assume there's a parent
             # transaction for this span that would be flushed out eventually.
             return None
-        
+
         print("")
 
         client = hub.client
@@ -333,7 +333,7 @@ def finish(self, hub=None):
             # We have no client and therefore nowhere to send this transaction
             # event.
             return None
-        
+
         print("LETS client")
 
         if not self.sampled:
diff --git a/tests/integrations/spark/test_spark.py b/tests/integrations/spark/test_spark.py
@@ -33,8 +33,10 @@ def test_set_app_properties():
         == sparkContext.applicationId
     )
 
+
 # SENTRY LISTENER
 
+
 def test_start_sentry_listener():
     sparkContext = SparkContext.getOrCreate()
 
@@ -197,8 +199,10 @@ def test_sentry_listener_on_stage_completed_failure(
     assert mockHub.kwargs["data"]["name"] == "run-job"
     assert mockHub.kwargs["data"]["reason"] == "failure-reason"
 
+
 # SPARK RDD
 
+
 def test_sentry_patch_rdd(sentry_init, capture_events, capture_exceptions):
     from pyspark.sql import SparkSession
     from operator import add
@@ -208,14 +212,13 @@ def test_sentry_patch_rdd(sentry_init, capture_events, capture_exceptions):
     events = capture_events()
     exceptions = capture_exceptions()
 
-    spark = SparkSession\
-        .builder\
-        .appName("PythonPi")\
-        .getOrCreate()
-    
+    spark = SparkSession.builder.appName("PythonPi").getOrCreate()
+
     spark.sparkContext.parallelize(range(1, 10), 2).map(lambda x: x + 2).reduce(add)
 
-    import pdb; pdb.set_trace()
+    import pdb
+
+    pdb.set_trace()
 
 
 ################
