Skip to content

Commit a9eab97

Browse files
committed
Simplify encoding kw_only args and test init=False args.
1. Replace the helper function for determining whether to use kw_only with a simpler to follow condition. 2. Add test coverage for the init=False case.
1 parent 1cadc5f commit a9eab97

File tree

2 files changed

+13
-15
lines changed

2 files changed

+13
-15
lines changed

sdks/python/apache_beam/coders/coder_impl.py

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232

3333
import decimal
3434
import enum
35-
import functools
3635
import itertools
3736
import json
3837
import logging
@@ -376,18 +375,6 @@ def _verify_dill_compat():
376375
raise RuntimeError(base_error + f". Found dill version '{dill.__version__}")
377376

378377

379-
dataclass_uses_kw_only: Callable[[Any], bool]
380-
if dataclasses:
381-
# Cache the result to avoid multiple checks for the same dataclass type.
382-
@functools.cache
383-
def dataclass_uses_kw_only(cls) -> bool:
384-
return any(
385-
field.init and field.kw_only for field in dataclasses.fields(cls))
386-
387-
else:
388-
dataclass_uses_kw_only = lambda cls: False
389-
390-
391378
class FastPrimitivesCoderImpl(StreamCoderImpl):
392379
"""For internal use only; no backwards-compatibility guarantees."""
393380
def __init__(
@@ -518,7 +505,7 @@ def encode_special_deterministic(self, value, stream):
518505
(value, type(value), self.requires_deterministic_step_label))
519506
init_fields = [field for field in dataclasses.fields(value) if field.init]
520507
try:
521-
if dataclass_uses_kw_only(type(value)):
508+
if any(field.kw_only for field in init_fields):
522509
stream.write_byte(DATACLASS_KW_ONLY_TYPE)
523510
self.encode_type(type(value), stream)
524511
stream.write_var_int64(len(init_fields))

sdks/python/apache_beam/coders/coders_test_common.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,15 @@ class UnFrozenDataClass:
123123
x: int
124124
y: int
125125

126+
@dataclasses.dataclass(frozen=True, kw_only=True)
127+
class FrozenUnInitKwOnlyDataClass:
128+
side: int
129+
area: int = dataclasses.field(init=False)
130+
131+
def __post_init__(self):
132+
# Hack to update an attribute in a frozen dataclass.
133+
object.__setattr__(self, 'area', self.side ** 2)
134+
126135

127136
# These tests need to all be run in the same process due to the asserts
128137
# in tearDownClass.
@@ -309,6 +318,7 @@ def test_deterministic_coder(self, compat_version):
309318
if dataclasses is not None:
310319
self.check_coder(deterministic_coder, FrozenDataClass(1, 2))
311320
self.check_coder(deterministic_coder, FrozenKwOnlyDataClass(c=1, d=2))
321+
self.check_coder(deterministic_coder, FrozenUnInitKwOnlyDataClass(side=11))
312322

313323
with self.assertRaises(TypeError):
314324
self.check_coder(deterministic_coder, UnFrozenDataClass(1, 2))
@@ -750,6 +760,7 @@ def test_cross_process_encoding_of_special_types_is_deterministic(
750760
from apache_beam.coders.coders_test_common import DefinesGetAndSetState
751761
from apache_beam.coders.coders_test_common import FrozenDataClass
752762
from apache_beam.coders.coders_test_common import FrozenKwOnlyDataClass
763+
from apache_beam.coders.coders_test_common import FrozenUnInitKwOnlyDataClass
753764
754765
755766
from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message
@@ -786,7 +797,7 @@ def test_cross_process_encoding_of_special_types_is_deterministic(
786797
("frozen_dataclass", FrozenDataClass(1, 2)),
787798
("frozen_dataclass_list", [FrozenDataClass(1, 2), FrozenDataClass(3, 4)]),
788799
("frozen_kwonly_dataclass", FrozenKwOnlyDataClass(c=1, d=2)),
789-
("frozen_kwonly_dataclass_list", [FrozenKwOnlyDataClass(c=1, d=2), FrozenKwOnlyDataClass(c=3, d=4)]),
800+
("frozen_kwonly_dataclass_list", [FrozenKwOnlyDataClass(c=1, d=2), FrozenUnInitKwOnlyDataClass(side=3)]),
790801
])
791802
792803
compat_version = {'"'+ compat_version +'"' if compat_version else None}

0 commit comments

Comments
 (0)