diff --git a/stdlib/src/builtin/builtin_list.mojo b/stdlib/src/builtin/builtin_list.mojo index 9e6b5972ea6..0b29b74110c 100644 --- a/stdlib/src/builtin/builtin_list.mojo +++ b/stdlib/src/builtin/builtin_list.mojo @@ -118,6 +118,33 @@ struct ListLiteral[*Ts: CollectionElement](Sized, CollectionElement): """ return value in self.storage + @always_inline + fn __getitem__[ + i: Int + ](ref self) -> ref [self.storage] __type_of(self.storage).element_types[ + i.value + ]: + """Get a list element at the given index. + + Parameters: + i: The element index. + + Returns: + The element at the given index. + + For example: + ```mojo + def main(): + x = [0, 1.0, "two"] + print( + x[0] == 0, + x[1] == 1.0, + x[2] == "two" + ) + ```. + """ + return self.storage[i] + # ===-----------------------------------------------------------------------===# # VariadicList / VariadicListMem diff --git a/stdlib/test/builtin/test_list_literal.mojo b/stdlib/test/builtin/test_list_literal.mojo index 309a831b241..50285cc0b92 100644 --- a/stdlib/test/builtin/test_list_literal.mojo +++ b/stdlib/test/builtin/test_list_literal.mojo @@ -12,7 +12,8 @@ # ===----------------------------------------------------------------------=== # # RUN: %mojo %s -from testing import assert_equal, assert_false, assert_true +from testing import assert_equal, assert_false, assert_true, assert_not_equal +from sys.intrinsics import _type_is_eq def test_list(): @@ -66,7 +67,25 @@ def test_contains(): assert_true("Mojo" not in h and String("Mojo") in h) +def test_list_literal_dunder_getitem(): + x = [0, True, 1.0, "two"] + assert_equal(x[0], 0) + assert_not_equal(x[0], 1) + assert_equal(x[1], True) + assert_not_equal(x[1], False) + assert_equal(x[2], 1.0) + assert_not_equal(x[2], 1.5) + assert_equal(x[3], "two") + assert_not_equal(x[3], "abc") + + assert_true(_type_is_eq[__type_of(x[0]), Int]()) + assert_true(_type_is_eq[__type_of(x[1]), Bool]()) + assert_true(_type_is_eq[__type_of(x[2]), Float64]()) + assert_true(_type_is_eq[__type_of(x[3]), StringLiteral]()) + + def main(): test_list() test_variadic_list() test_contains() + test_list_literal_dunder_getitem()