fix formatter and wrong test type

Author: ian-Liaozy

sdks/python/apache_beam/runners/interactive/interactive_beam.py

@@ -1056,8 +1056,7 @@ def compute(
       except TypeError:
         raise ValueError(
             f'The given pcoll {pcoll_container} is not a dict, an iterable or '
-            'a PCollection.'
-        )
+            'a PCollection.')
 
   pcolls_set = set()
   for pcoll in flatten_pcolls:
@@ -1084,8 +1083,7 @@ def compute(
       raise ValueError('All PCollections must belong to the same pipeline.')
 
   recording_manager = ie.current_env().get_recording_manager(
-      user_pipeline, create_if_absent=True
-  )
+      user_pipeline, create_if_absent=True)
 
   return recording_manager.compute_async(
       pcolls_set,

sdks/python/apache_beam/runners/interactive/interactive_beam_test.py

@@ -675,7 +675,6 @@ def test_default_value_for_invalid_worker_number(self):
 
 @isolated_env
 class InteractiveBeamComputeTest(unittest.TestCase):
-
   def setUp(self):
     self.env = ie.current_env()
     self.env._is_in_ipython = False  # Default to non-IPython
@@ -820,9 +819,9 @@ def test_compute_dependency_wait_true(self):
     self.assertTrue(self.env.is_pcollection_computing(pcoll1))
 
     # Spy on _wait_for_dependencies
-    with patch.object(
-        rm, '_wait_for_dependencies', wraps=rm._wait_for_dependencies
-    ) as spy_wait:
+    with patch.object(rm,
+                      '_wait_for_dependencies',
+                      wraps=rm._wait_for_dependencies) as spy_wait:
       async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=True)
 
       # Check that wait_for_dependencies was called for pcoll2
@@ -850,9 +849,9 @@ def test_compute_dependency_wait_false(self, mock_is_computing):
     # Pretend pcoll1 is computing
     mock_is_computing.side_effect = lambda pcoll: pcoll is pcoll1
 
-    with patch.object(
-        rm, '_execute_pipeline_fragment', wraps=rm._execute_pipeline_fragment
-    ) as spy_execute:
+    with patch.object(rm,
+                      '_execute_pipeline_fragment',
+                      wraps=rm._execute_pipeline_fragment) as spy_execute:
       async_res2 = ib.compute(pcoll2, blocking=False, wait_for_inputs=False)
       async_res2.result(timeout=60)
 

sdks/python/apache_beam/runners/interactive/recording_manager.py

@@ -64,7 +64,6 @@
 
 class AsyncComputationResult:
   """Represents the result of an asynchronous computation."""
-
   def __init__(
       self,
       future: Future,
@@ -81,19 +80,15 @@ def __init__(
     self._display_id = str(uuid.uuid4())
     self._output_widget = widgets.Output() if IS_IPYTHON else None
     self._cancel_button = (
-        widgets.Button(description='Cancel') if IS_IPYTHON else None
-    )
+        widgets.Button(description='Cancel') if IS_IPYTHON else None)
     self._progress_bar = (
         widgets.FloatProgress(
             value=0.0,
             min=0.0,
             max=1.0,
             description='Running:',
             bar_style='info',
-        )
-        if IS_IPYTHON
-        else None
-    )
+        ) if IS_IPYTHON else None)
     self._cancel_requested = False
 
     if IS_IPYTHON:
@@ -143,8 +138,7 @@ def update_display(self, msg: str, progress: Optional[float] = None):
       self._cancel_button.disabled = False
 
   def set_pipeline_result(
-      self, pipeline_result: beam.runners.runner.PipelineResult
-  ):
+      self, pipeline_result: beam.runners.runner.PipelineResult):
     self._pipeline_result = pipeline_result
     if self._cancel_requested:
       self.cancel()
@@ -181,8 +175,7 @@ def _on_done(self, future: Future):
       else:
         _LOGGER.warning(
             'Async computation finished but state is not DONE:'
-            f" {res.state if res else 'Unknown'}"
-        )
+            f" {res.state if res else 'Unknown'}")
 
   def cancel(self):
     if self._future.done():
@@ -200,8 +193,7 @@ def cancel(self):
         if PipelineState.is_terminal(current_state):
           self.update_display(
               'Cannot cancel: Pipeline already in terminal state'
-              f' {current_state}.'
-          )
+              f' {current_state}.')
           return False
 
         self._pipeline_result.cancel()
@@ -221,8 +213,7 @@ def __repr__(self):
     return (
         f'<AsyncComputationResult({self._display_id}) for'
         f' {len(self._pcolls)} PCollections, status:'
-        f" {'done' if self.done() else 'running'}>"
-    )
+        f" {'done' if self.done() else 'running'}>")
 
 
 class ElementStream:
@@ -475,7 +466,7 @@ def _execute_pipeline_fragment(
     })
 
     fragment = pf.PipelineFragment(
-      list(pcolls_to_compute), merged_options, runner=runner)
+        list(pcolls_to_compute), merged_options, runner=runner)
 
     if async_result:
       async_result.update_display('Building pipeline fragment...', 0.1)
@@ -508,12 +499,10 @@ def _run_async_computation(
       self._env.mark_pcollection_computing(pcolls_to_compute)
       _LOGGER.info(
           'Starting asynchronous computation for'
-          f' {len(pcolls_to_compute)} PCollections.'
-      )
+          f' {len(pcolls_to_compute)} PCollections.')
 
       pipeline_result = self._execute_pipeline_fragment(
-          pcolls_to_compute, async_result, runner, options
-      )
+          pcolls_to_compute, async_result, runner, options)
 
       # if pipeline_result.state == PipelineState.DONE:
       #   self._env.mark_pcollection_computed(pcolls_to_compute)
@@ -668,18 +657,19 @@ def compute_async(
       self._env.evict_computed_pcollections(self.user_pipeline)
 
     computed_pcolls = {
-        pcoll for pcoll in pcolls if pcoll in self._env.computed_pcollections
+        pcoll
+        for pcoll in pcolls if pcoll in self._env.computed_pcollections
     }
     computing_pcolls = {
-        pcoll for pcoll in pcolls if self._env.is_pcollection_computing(pcoll)
+        pcoll
+        for pcoll in pcolls if self._env.is_pcollection_computing(pcoll)
     }
     pcolls_to_compute = pcolls - computed_pcolls - computing_pcolls
 
     if not pcolls_to_compute:
       _LOGGER.info(
           'All requested PCollections are already computed or are being'
-          ' computed.'
-      )
+          ' computed.')
       return None
 
     self._watch(list(pcolls_to_compute))
@@ -691,36 +681,30 @@ def compute_async(
         if wait_for_inputs:
           if not self._wait_for_dependencies(pcolls_to_compute):
             raise RuntimeError(
-                'Dependency computation failed or was cancelled.'
-            )
+                'Dependency computation failed or was cancelled.')
         pipeline_result = self._execute_pipeline_fragment(
-            pcolls_to_compute, None, runner, options
-        )
+            pcolls_to_compute, None, runner, options)
         if pipeline_result.state == PipelineState.DONE:
           self._env.mark_pcollection_computed(pcolls_to_compute)
         else:
           _LOGGER.error(
-              f'Blocking computation failed. State: {pipeline_result.state}'
-          )
+              f'Blocking computation failed. State: {pipeline_result.state}')
           raise RuntimeError(
-              f'Blocking computation failed. State: {pipeline_result.state}'
-          )
+              f'Blocking computation failed. State: {pipeline_result.state}')
       finally:
         self._env.unmark_pcollection_computing(pcolls_to_compute)
       return None
 
     else:  # Asynchronous
       future = Future()
       async_result = AsyncComputationResult(
-          future, pcolls_to_compute, self.user_pipeline, self
-      )
+          future, pcolls_to_compute, self.user_pipeline, self)
       self._async_computations[async_result._display_id] = async_result
 
       def task():
         try:
           result = self._run_async_computation(
-              pcolls_to_compute, async_result, wait_for_inputs, runner, options
-          )
+              pcolls_to_compute, async_result, wait_for_inputs, runner, options)
           future.set_result(result)
         except Exception as e:
           if not future.cancelled():
@@ -744,8 +728,8 @@ def _get_pcoll_id_map(self):
     return {v: k for k, v in pcoll_to_id.items()}
 
   def _get_all_dependencies(
-      self, pcolls: Set[beam.pvalue.PCollection]
-  ) -> Set[beam.pvalue.PCollection]:
+      self,
+      pcolls: Set[beam.pvalue.PCollection]) -> Set[beam.pvalue.PCollection]:
     """Gets all upstream PCollection dependencies for the given set of PCollections."""
     if not self._pipeline_graph:
       return set()
@@ -759,8 +743,7 @@ def _get_all_dependencies(
 
     target_pcoll_ids = {
         pcoll_to_id.get(str(pcoll))
-        for pcoll in pcolls
-        if str(pcoll) in pcoll_to_id
+        for pcoll in pcolls if str(pcoll) in pcoll_to_id
     }
 
     if not target_pcoll_ids:
@@ -824,16 +807,13 @@ def _wait_for_dependencies(
 
     if async_result:
       async_result.update_display(
-          f'Waiting for {len(computing_deps)} dependencies to finish...'
-      )
+          f'Waiting for {len(computing_deps)} dependencies to finish...')
     _LOGGER.info(
         f'Waiting for {len(computing_deps)} dependencies:'
-        f' {computing_deps.keys()}'
-    )
+        f' {computing_deps.keys()}')
 
     futures_to_wait = list(
-        set(comp._future for comp in computing_deps.values())
-    )
+        set(comp._future for comp in computing_deps.values()))
 
     try:
       for i, future in enumerate(futures_to_wait):
@@ -901,8 +881,7 @@ def record(
       if not self._wait_for_dependencies(uncomputed_pcolls):
         raise RuntimeError(
             'Cannot record because a dependency failed to compute'
-            ' asynchronously.'
-        )
+            ' asynchronously.')
 
       self._clear()
 

sdks/python/apache_beam/runners/interactive/recording_manager_test.py

@@ -45,8 +45,7 @@
 from apache_beam.utils.windowed_value import WindowedValue
 
 
-class AsyncComputationResultTest(googletest.TestCase):
-
+class AsyncComputationResultTest(unittest.TestCase):
   def setUp(self):
     self.mock_future = MagicMock(spec=Future)
     self.pcolls = {MagicMock(spec=beam.pvalue.PCollection)}