From f1204d08bf21484e030480d981d2cfeb55e12301 Mon Sep 17 00:00:00 2001
From: Aaron Berdy <berdy@amazon.com>
Date: Mon, 20 Nov 2023 18:37:01 -0800
Subject: [PATCH] feat: add str, repr and getitem to BasisState

---
 src/braket/circuits/basis_state.py            |  9 ++++
 .../braket/circuits/test_basis_state.py       | 46 +++++++++++++++++--
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/src/braket/circuits/basis_state.py b/src/braket/circuits/basis_state.py
index 814444e75..825a074fb 100644
--- a/src/braket/circuits/basis_state.py
+++ b/src/braket/circuits/basis_state.py
@@ -33,6 +33,15 @@ def __iter__(self):
     def __eq__(self, other):
         return self.state == other.state
 
+    def __str__(self):
+        return self.as_string
+
+    def __repr__(self):
+        return f'BasisState("{self.as_string}")'
+
+    def __getitem__(self, item):
+        return BasisState(self.state[item])
+
 
 BasisStateInput = Union[int, list[int], str, BasisState]
 
diff --git a/test/unit_tests/braket/circuits/test_basis_state.py b/test/unit_tests/braket/circuits/test_basis_state.py
index 023494fae..3e08e45a6 100644
--- a/test/unit_tests/braket/circuits/test_basis_state.py
+++ b/test/unit_tests/braket/circuits/test_basis_state.py
@@ -51,6 +51,46 @@
     ),
 )
 def test_as_props(basis_state_input, size, as_tuple, as_int, as_string):
-    assert BasisState(basis_state_input, size).as_tuple == as_tuple
-    assert BasisState(basis_state_input, size).as_int == as_int
-    assert BasisState(basis_state_input, size).as_string == as_string
+    basis_state = BasisState(basis_state_input, size)
+    assert basis_state.as_tuple == as_tuple
+    assert basis_state.as_int == as_int
+    assert basis_state.as_string == as_string == str(basis_state)
+
+
+@pytest.mark.parametrize(
+    "basis_state_input, index, substate_input",
+    (
+        (
+            "1001",
+            slice(None),
+            "1001",
+        ),
+        (
+            "1001",
+            3,
+            "1",
+        ),
+        (
+            "1010",
+            slice(None, None, 2),
+            "11",
+        ),
+        (
+            "1010",
+            slice(1, None, 2),
+            "00",
+        ),
+        (
+            "1010",
+            slice(None, -2),
+            "10",
+        ),
+        (
+            "1010",
+            -1,
+            "0",
+        ),
+    ),
+)
+def test_indexing(basis_state_input, index, substate_input):
+    assert BasisState(basis_state_input)[index] == BasisState(substate_input)