Add type checking to confirm that the flatten inputs are actually pco… (#35874)

* Add type checking to confirm that the flatten inputs are actually pcollections

* Allow iterables of non-pcollections to flatten again

* Update sdks/python/apache_beam/transforms/core.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: hjtran

sdks/python/apache_beam/transforms/core.py

@@ -3866,6 +3866,15 @@ def _extract_input_pvalues(self, pvalueish):
       raise ValueError(
           'Input to Flatten must be an iterable. '
           'Got a value of type %s instead.' % type(pvalueish))
+
+    # Spot check to see if any of the items are iterables of PCollections
+    # and raise an error if so. This is always a user-error
+    for idx, item in enumerate(pvalueish):
+      if isinstance(item, (list, tuple)) and any(
+          isinstance(sub_item, pvalue.PCollection) for sub_item in item):
+        raise TypeError(
+            'Inputs to Flatten cannot include an iterable of PCollections. '
+            f'(input at index {idx}: "{item}")')
     return pvalueish, pvalueish
 
   def expand(self, pcolls):

sdks/python/apache_beam/typehints/typecheck_test.py

@@ -179,6 +179,15 @@ def process(self, element, *args, **kwargs):
       (self.p | beam.Create(['1', '1']) | beam.ParDo(ToInt()))
       self.p.run().wait_until_finish()
 
+  def test_bad_flatten_input(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        "Inputs to Flatten cannot include an iterable of PCollections. "):
+      with beam.Pipeline() as p:
+        pc = p | beam.Create([1, 1])
+        flatten_inputs = [pc, (pc, )]
+        flatten_inputs | beam.Flatten()
+
   def test_do_fn_returning_non_iterable_throws_error(self):
     # This function is incorrect because it returns a non-iterable object
     def incorrect_par_do_fn(x):