Repository: quora/pyanalyze
Pull request #462 (2022-02-06 02:24:29)

fix getitem on list/tuple subclasses


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 4bf6dab..cd5dd19 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Fix incorrect type inferred for indexing operations on
+  subclasses of `list` and `tuple` (#461)
 - Add plugin providing a precise type for `dict.get` calls (#460)
 - Fix internal error when an `__eq__` method throws (#461)
 - Fix handling of `async def` methods in stubs (#459)
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 003575a..d701e06 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -398,8 +398,16 @@ def _sequence_getitem_impl(ctx: CallContext, typ: type) -> ImplReturn:
                     return SequenceIncompleteValue.make_or_known(
                         list, self_value.members[key.val]
                     )
-                else:
+                elif self_value.typ in (list, tuple):
+                    # For generics of exactly list/tuple, return the self type.
                     return self_value
+                else:
+                    # slicing a subclass of list or tuple returns a list
+                    # or tuple, not a subclass (unless the subclass overrides
+                    # __getitem__, but then we wouldn't get here).
+                    # TODO return a more precise type if the class inherits
+                    # from a generic list/tuple.
+                    return TypedValue(typ)
             else:
                 ctx.show_error(f"Invalid {typ.__name__} key {key}")
                 return AnyValue(AnySource.error)


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index dca9c5f..f8447a2 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -879,6 +879,7 @@ class TestSequenceGetItem(TestNameCheckVisitorBase):
             assert_is_value(tpl[0], TypedValue(int))
             assert_is_value(tpl[-1], TypedValue(int))
             assert_is_value(tpl[:1], GenericValue(tuple, [TypedValue(int)]))
+            assert_is_value(tpl[:], GenericValue(tuple, [TypedValue(int)]))
             assert_is_value(tpl[i], TypedValue(int))
             assert_is_value(tpl[s], GenericValue(tuple, [TypedValue(int)]))
             assert_is_value(tpl[unannotated], AnyValue(AnySource.from_another))
@@ -939,6 +940,26 @@ class TestSequenceGetItem(TestNameCheckVisitorBase):
             sorted_indexes = sorted(range(len(words)), key=lambda i: words[i])
             return sorted_indexes
 
+    @assert_passes()
+    def test_subclasses(self):
+        import time
+
+        class MyList(list):
+            pass
+
+        class MyTuple(tuple):
+            pass
+
+        def capybara(t: time.struct_time, ml: MyList, mt: MyTuple):
+            assert_is_value(t[0], TypedValue(int))
+            assert_is_value(t[:], TypedValue(tuple))
+            assert_is_value(t[:6], TypedValue(tuple))
+
+            assert_is_value(ml[0], AnyValue(AnySource.generic_argument))
+            assert_is_value(ml[:], TypedValue(list))
+            assert_is_value(mt[0], AnyValue(AnySource.generic_argument))
+            assert_is_value(mt[:], TypedValue(tuple))
+
 
 class TestDictGetItem(TestNameCheckVisitorBase):
     @assert_fails(ErrorCode.unhashable_key)