diff --git a/skoolkit/skoolmacro.py b/skoolkit/skoolmacro.py index 88d5f2c3..94e17784 100644 --- a/skoolkit/skoolmacro.py +++ b/skoolkit/skoolmacro.py @@ -943,8 +943,18 @@ def parse_foreach(writer, text, index, *cwd): values = [str(e.address) for e in entry_holder.memory_map if e.ctl != 'i' and (not types or e.ctl in types)] elif value.startswith('POKE'): name = value[4:] + m = re.match(r'(.*)\[([0-9]+|[0-9]*:[0-9]*)\]$', name) + if m: + name = m.group(1) + elements = [e or (0, None)[i] for i, e in enumerate(m.group(2).split(':', 1))] + if len(elements) == 1: + indexes = [int(elements[0]), int(elements[0]) + 1] + else: + indexes = [int(e) if e else None for e in elements] + else: + indexes = (None, None) values = [] - for addr, byte, length, step in writer.pokes[name]: + for addr, byte, length, step in writer.pokes[name][slice(*indexes)]: if length == 1: values.append(f'POKE {addr},{byte}') else: diff --git a/sphinx/source/skool-macros.rst b/sphinx/source/skool-macros.rst index c8e4dbf1..e28ac696 100644 --- a/sphinx/source/skool-macros.rst +++ b/sphinx/source/skool-macros.rst @@ -481,6 +481,16 @@ example: * ``#POKES30000,1,2`` gives 'FOR n=30000 TO 30001: POKE n,1: NEXT n' * ``#POKES30000,1,2,3`` gives 'FOR n=30000 TO 30003 STEP 3: POKE n,1: NEXT n' +To specify a subset of the POKEs in a list, use either an index parameter, or +start and stop parameters separated by a colon, enclosed in square brackets and +appended to ``name``. As with Python sequences, these parameters are 0-based, +and the start and stop parameters are optional, so: + +* ``POKEname[1]`` yields the second POKE in the list +* ``POKEname[1:3]`` yields the second and third POKEs in the list +* ``POKEname[1:]`` yields all POKEs in the list except the first +* ``POKEname[:3]`` yields the first three POKEs in the list + See :ref:`stringParameters` for details on alternative ways to supply the ``s1,s2,...`` and ``var,string[,sep,fsep]`` parameter strings. diff --git a/tests/macrotest.py b/tests/macrotest.py index 6ba09c9c..402b9316 100644 --- a/tests/macrotest.py +++ b/tests/macrotest.py @@ -1054,6 +1054,31 @@ def test_macro_foreach_with_poke_two_snapshots(self): writer.expand('#PUSHSthud #POKES30000,10 #POPS') self.assertEqual(writer.expand('#FOREACH(POKEthud)(a,a)'), 'POKE 30000,10') + def test_macro_foreach_with_poke_indexing(self): + skool = """ + @start + b30000 DEFB 1,1,1 + """ + writer = self._get_writer(skool=skool) + writer.expand('#PUSHSaaa #POKES30000,11;30001,12;30002,13 #POPS') + self.assertEqual(writer.expand('#FOREACH(POKEaaa[0])(p,p)'), 'POKE 30000,11') + self.assertEqual(writer.expand('#FOREACH(POKEaaa[1])(p,p)'), 'POKE 30001,12') + self.assertEqual(writer.expand('#FOREACH(POKEaaa[2])(p,p)'), 'POKE 30002,13') + + def test_macro_foreach_with_poke_slicing(self): + skool = """ + @start + b30000 DEFB 0,0,0 + """ + writer = self._get_writer(skool=skool) + writer.expand('#PUSHSbbb #POKES30000,1;30001,2;30002,3 #POPS') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[0:1])(p,p)'), 'POKE 30000,1') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[0:2])(p,[p])'), '[POKE 30000,1][POKE 30001,2]') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[:2])(p,[p])'), '[POKE 30000,1][POKE 30001,2]') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[1:3])(p,[p])'), '[POKE 30001,2][POKE 30002,3]') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[1:])(p,[p])'), '[POKE 30001,2][POKE 30002,3]') + self.assertEqual(writer.expand('#FOREACH(POKEbbb[:])(p,[p])'), '[POKE 30000,1][POKE 30001,2][POKE 30002,3]') + def test_macro_foreach_invalid(self): writer = self._get_writer() prefix = ERROR_PREFIX.format('FOREACH')