From e73f5812803999440093dead8af6bf66b11bfd6f Mon Sep 17 00:00:00 2001 From: Robert Bradshaw Date: Thu, 5 Sep 2019 10:40:57 -0700 Subject: [PATCH] Fix slice inference to not always assume List. --- sdks/python/apache_beam/typehints/opcodes.py | 4 ++-- .../python/apache_beam/typehints/trivial_inference_test.py | 7 +++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/typehints/opcodes.py b/sdks/python/apache_beam/typehints/opcodes.py index eb33ada984f0..6aa5de21168c 100644 --- a/sdks/python/apache_beam/typehints/opcodes.py +++ b/sdks/python/apache_beam/typehints/opcodes.py @@ -161,8 +161,8 @@ def binary_subscr(state, unused_arg): out = base._constraint_for_index(index.value) except IndexError: out = element_type(base) - elif index == slice: - out = typehints.List[element_type(base)] + elif index == slice and isinstance(base, typehints.IndexableTypeConstraint): + out = base else: out = element_type(base) state.stack.append(out) diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py b/sdks/python/apache_beam/typehints/trivial_inference_test.py index e32a375f27d6..ff0949bbd57a 100644 --- a/sdks/python/apache_beam/typehints/trivial_inference_test.py +++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py @@ -60,8 +60,15 @@ def reverse(ab): typehints.Tuple[int, str], reverse, [typehints.Tuple[str, float, int]]) self.assertReturnType( typehints.Tuple[int, int], reverse, [typehints.List[int]]) + + def testGetItemSlice(self): self.assertReturnType( typehints.List[int], lambda v: v[::-1], [typehints.List[int]]) + self.assertReturnType( + typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]]) + self.assertReturnType(str, lambda v: v[::-1], [str]) + self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any]) + self.assertReturnType(typehints.Any, lambda v: v[::-1], [object]) def testUnpack(self): def reverse(a_b):