Skip to content

Commit

Permalink
[SPARK-44833][CONNECT][PYTHON][FOLLOW-UP] Fix the type annotation for…
Browse files Browse the repository at this point in the history
… iterator at reattach.py

### What changes were proposed in this pull request?

This PR fixes the type hint of `self._iterator` at `reattach.py` to be `Optional` because `None` is assigned here.

### Why are the changes needed?

To fix the CI:

```
starting mypy annotations test...
annotations failed mypy checks:
python/pyspark/sql/connect/client/reattach.py:149: error: Incompatible types in assignment (expression has type "None", variable has type "Iterator[ExecutePlanResponse]")  [assignment]
python/pyspark/sql/connect/client/reattach.py:254: error: Incompatible types in assignment (expression has type "None", variable has type "Iterator[ExecutePlanResponse]")  [assignment]
python/pyspark/sql/connect/client/reattach.py:[25](https://github.com/apache/spark/actions/runs/6093065151/job/16532169212#step:19:26)8: error: Incompatible types in assignment (expression has type "None", variable has type "Iterator[ExecutePlanResponse]")  [assignment]
Found 3 errors in 1 file (checked 703 source files)
```

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Manually ran the script of `./dev/lint-python`

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #42830 from HyukjinKwon/SPARK-44833-followup.

Authored-by: Hyukjin Kwon <gurwls223@apache.org>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
(cherry picked from commit 48872d4)
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
  • Loading branch information
HyukjinKwon committed Sep 6, 2023
1 parent 9b750e9 commit c98ade3
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions python/pyspark/sql/connect/client/reattach.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def __init__(
# Note: This is not retried, because no error would ever be thrown here, and GRPC will only
# throw error on first self._has_next().
self._metadata = metadata
self._iterator: Iterator[pb2.ExecutePlanResponse] = iter(
self._iterator: Optional[Iterator[pb2.ExecutePlanResponse]] = iter(
self._stub.ExecutePlan(self._initial_request, metadata=metadata)
)

Expand Down Expand Up @@ -133,7 +133,9 @@ def _has_next(self) -> bool:
with attempt:
if self._current is None:
try:
self._current = self._call_iter(lambda: next(self._iterator))
self._current = self._call_iter(
lambda: next(self._iterator) # type: ignore[arg-type]
)
except StopIteration:
pass

Expand All @@ -150,7 +152,9 @@ def _has_next(self) -> bool:
# shouldn't change
assert not self._result_complete
try:
self._current = self._call_iter(lambda: next(self._iterator))
self._current = self._call_iter(
lambda: next(self._iterator) # type: ignore[arg-type]
)
except StopIteration:
pass
has_next = self._current is not None
Expand Down

0 comments on commit c98ade3

Please sign in to comment.