Re: [PR] Add validation to check for non-ints from partitionfns [beam]

[via GitHub]

gemini-code-assist[bot] commented on code in PR #35881:
URL: https://github.com/apache/beam/pull/35881#discussion_r2280439205


##########
sdks/python/apache_beam/transforms/core_test.py:
##########
@@ -162,6 +162,14 @@ def 
test_dofn_with_implicit_return_none_return_without_value(self):
 
 
 class PartitionTest(unittest.TestCase):
+  def test_partition_with_bools(self):
+    with pytest.raises(
+        Exception,
+        match="PartitionFn yielded a 'bool' when it should only yields 
integers"
+    ):

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The test correctly checks for the exception, but it's good practice to be 
more specific. The code raises a `ValueError`, so the test should expect a 
`ValueError` instead of a generic `Exception`. This makes the test more robust. 
Also, there's a small grammatical error in the match string; it should be 
'yield' instead of 'yields'.
   
   ```suggestion
       with pytest.raises(
           ValueError,
           match="PartitionFn yielded a 'bool' when it should only yield 
integers"
       ):
   ```



##########
sdks/python/apache_beam/transforms/core.py:
##########
@@ -3573,6 +3573,10 @@ class ApplyPartitionFnFn(DoFn):
     """A DoFn that applies a PartitionFn."""
     def process(self, element, partitionfn, n, *args, **kwargs):
       partition = partitionfn.partition_for(element, n, *args, **kwargs)
+      if isinstance(partition, bool) or not isinstance(partition, int):
+        raise ValueError(
+            f"PartitionFn yielded a '{type(partition).__name__}' "
+            "when it should only yields integers")

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The condition `isinstance(partition, bool) or not isinstance(partition, 
int)` is correct, but it can be simplified. Since `bool` is a subclass of 
`int`, `isinstance(True, int)` evaluates to `True`. A more direct way to check 
if a value is strictly an integer (and not a subclass like `bool`) is to use 
`type(partition) is not int`. This is more readable and achieves the same 
result. I've also corrected a small grammatical error in the error message 
('yields' to 'yield').
   
   ```suggestion
         if type(partition) is not int:
           raise ValueError(
               f"PartitionFn yielded a '{type(partition).__name__}' "
               "when it should only yield integers")
   ```



##########
sdks/python/apache_beam/transforms/core_test.py:
##########
@@ -162,6 +162,14 @@
 
 
 class PartitionTest(unittest.TestCase):
+  def test_partition_with_bools(self):
+    with pytest.raises(
+        Exception,
+        match="PartitionFn yielded a 'bool' when it should only yields 
integers"
+    ):
+      with beam.testing.test_pipeline.TestPipeline() as p:
+        _ = (p | beam.Create([True]) | beam.Partition(lambda x, _: x, 2))

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   This test covers the case where the partition function returns a boolean. To 
make the tests more comprehensive, it would be beneficial to add another test 
case for other non-integer types that should be rejected, such as floats. This 
would ensure the validation logic is robust against various invalid types.
   
   For example, you could add a test like this:
   ```python
     def test_partition_with_floats(self):
       with pytest.raises(
           ValueError,
           match="PartitionFn yielded a 'float' when it should only yield 
integers"
       ):
         with beam.testing.test_pipeline.TestPipeline() as p:
           _ = (p | beam.Create([1.0]) | beam.Partition(lambda x, _: x, 2))
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@beam.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org