From ecc5490f11f67e6a7d6b851db7aecbe10b063902 Mon Sep 17 00:00:00 2001
From: rd4com <144297616+rd4com@users.noreply.github.com>
Date: Wed, 4 Dec 2024 20:29:48 +0100
Subject: [PATCH] [stdlib] Add `ListLiteral.__getitem__`

Signed-off-by: rd4com <144297616+rd4com@users.noreply.github.com>
---
 stdlib/src/builtin/builtin_list.mojo       | 27 ++++++++++++++++++++++
 stdlib/test/builtin/test_list_literal.mojo | 21 ++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

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()