Address feedback from PR #241 (#244)

I added a function called `Token.with_source` to simplify the
repeated logic in the source extent tests, but if it's unnecessary I can
easily revert those commits.

Edit: Linking the PR for easy, future reference: #241.

Author: neuroevolutus

scrapscript.py

@@ -3,6 +3,7 @@
 import argparse
 import base64
 import code
+import copy
 import dataclasses
 import enum
 import functools
@@ -15,7 +16,6 @@
 import urllib.request
 from dataclasses import dataclass
 from enum import auto
-from functools import reduce
 from types import ModuleType
 from typing import Any, Callable, Dict, Generator, Iterator, Mapping, Optional, Set, Tuple, Union
 
@@ -59,6 +59,10 @@ def join_source_extents(
 class Token:
     source_extent: SourceExtent = dataclasses.field(default_factory=SourceExtent, init=False, compare=False)
 
+    def with_source(self, source_extent: SourceExtent) -> Token:
+        self.source_extent = source_extent
+        return self
+
 
 @dataclass(eq=True)
 class IntLit(Token):
@@ -213,18 +217,7 @@ def peek_char(self) -> str:
 
     def make_token(self, cls: type, *args: Any) -> Token:
         result: Token = cls(*args)
-
-        # Set start of token's source extent
-        result.source_extent.start.lineno = self.current_token_source_extent.start.lineno
-        result.source_extent.start.colno = self.current_token_source_extent.start.colno
-        result.source_extent.start.byteno = self.current_token_source_extent.start.byteno
-
-        # Set end of token's source extent
-        result.source_extent.end.colno = self.current_token_source_extent.end.colno
-        result.source_extent.end.lineno = self.current_token_source_extent.end.lineno
-        result.source_extent.end.byteno = self.current_token_source_extent.end.byteno
-
-        return result
+        return result.with_source(copy.deepcopy(self.current_token_source_extent))
 
     def read_tokens(self) -> Generator[Token, None, None]:
         while (token := self.read_token()) and not isinstance(token, EOF):
@@ -568,6 +561,7 @@ def parse_unary(tokens: Peekable, p: float) -> "Object":
             MatchCase, pipe_source_extent.coalesce(expr.source_extent), expr.arg, expr.body
         )
         cases = [match_case]
+        match_function_source_extent = match_case.source_extent
         while True:
             try:
                 if tokens.peek() != Operator("|"):
@@ -582,13 +576,10 @@ def parse_unary(tokens: Peekable, p: float) -> "Object":
                 MatchCase, pipe_source_extent.coalesce(expr.source_extent), expr.arg, expr.body
             )
             cases.append(match_case)
-        cases_source_extents = [case_branch.source_extent for case_branch in cases]
+            match_function_source_extent = join_source_extents(match_function_source_extent, match_case.source_extent)
         return make_source_annotated_object(
             MatchFunction,
-            reduce(
-                join_source_extents,
-                cases_source_extents,
-            ),
+            match_function_source_extent,
             cases,
         )
     elif isinstance(token, LeftParen):

scrapscript_tests.py

@@ -1,4 +1,3 @@
-import base64
 import unittest
 import re
 from typing import Optional
@@ -1201,172 +1200,149 @@ def test_apply_with_variant_args(self) -> None:
         self.assertEqual(ast, Apply(Apply(Var("f"), TRUE), FALSE))
 
     def test_parse_int_preserves_source_extent(self) -> None:
-        int_lit = IntLit(1)
         source_extent = SourceExtent(
             start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
         )
-        int_lit.source_extent = source_extent
-        int_ast = make_source_annotated_object(Int, source_extent, 1)
-        self.assertTrue(parse(Peekable(iter([int_lit]))).source_extent == int_ast.source_extent)
+        int_lit = IntLit(1).with_source(source_extent)
+        self.assertEqual(parse(Peekable(iter([int_lit]))).source_extent, source_extent)
 
     def test_parse_float_preserves_source_extent(self) -> None:
-        float_lit = FloatLit(3.2)
         source_extent = SourceExtent(
             start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=3, byteno=2)
         )
-        float_lit.source_extent = source_extent
-        float_ast = make_source_annotated_object(Float, source_extent, 3.2)
-        self.assertTrue(parse(Peekable(iter([float_lit]))).source_extent == float_ast.source_extent)
+        float_lit = FloatLit(3.2).with_source(source_extent)
+        self.assertEqual(parse(Peekable(iter([float_lit]))).source_extent, source_extent)
 
     def test_parse_string_preserves_source_extent(self) -> None:
-        string_lit = StringLit("Hello")
         source_extent = SourceExtent(
             start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=7, byteno=6)
         )
-        string_lit.source_extent = source_extent
-        string_ast = make_source_annotated_object(String, source_extent, "Hello")
-        self.assertTrue(parse(Peekable(iter([string_lit]))).source_extent == string_ast.source_extent)
+        string_lit = StringLit("Hello").with_source(source_extent)
+        self.assertEqual(parse(Peekable(iter([string_lit]))).source_extent, source_extent)
 
     def test_parse_bytes_preserves_source_extent(self) -> None:
-        bytes_lit = BytesLit("QUJD", 64)
         source_extent = SourceExtent(
             start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=9, byteno=8)
         )
-        bytes_lit.source_extent = source_extent
-        bytes_ast = make_source_annotated_object(Bytes, source_extent, base64.b64decode("QUJD"))
-        self.assertTrue(parse(Peekable(iter([bytes_lit]))).source_extent == bytes_ast.source_extent)
+        bytes_lit = BytesLit("QUJD", 64).with_source(source_extent)
+        self.assertEqual(parse(Peekable(iter([bytes_lit]))).source_extent, source_extent)
 
     def test_parse_var_preserves_source_extent(self) -> None:
-        var = Name("x")
         source_extent = SourceExtent(
             start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
         )
-        var.source_extent = source_extent
-        var_ast = make_source_annotated_object(Var, source_extent, "x")
-        self.assertTrue(parse(Peekable(iter([var]))).source_extent == var_ast.source_extent)
+        var = Name("x").with_source(source_extent)
+        self.assertEqual(parse(Peekable(iter([var]))).source_extent, source_extent)
 
     def test_parse_hole_preserves_source_extent(self) -> None:
-        left_paren = LeftParen()
-        left_paren.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
-        )
-        right_paren = RightParen()
-        right_paren.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+        left_paren = LeftParen().with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
+            )
         )
-        hole = make_source_annotated_object(
-            Hole,
+        right_paren = RightParen().with_source(
             SourceExtent(
-                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=2, byteno=1)
-            ),
+                start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+            )
+        )
+        hole_source_extent = SourceExtent(
+            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=2, byteno=1)
         )
-        self.assertTrue(parse(Peekable(iter([left_paren, right_paren]))).source_extent == hole.source_extent)
+        self.assertEqual(parse(Peekable(iter([left_paren, right_paren]))).source_extent, hole_source_extent)
 
     def test_parenthesized_expression_preserves_source_extent(self) -> None:
-        left_paren = LeftParen()
-        left_paren.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
-        )
-        int_lit = IntLit(1)
-        int_lit.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+        left_paren = LeftParen().with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
+            )
         )
-        right_paren = RightParen()
-        right_paren.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+        int_lit = IntLit(1).with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+            )
         )
-        parenthesized_int_lit = make_source_annotated_object(
-            Int,
+        right_paren = RightParen().with_source(
             SourceExtent(
-                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=3, byteno=2)
-            ),
-            1,
+                start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+            )
         )
-        self.assertTrue(
-            parse(Peekable(iter([left_paren, int_lit, right_paren]))).source_extent
-            == parenthesized_int_lit.source_extent
+        parenthesized_int_lit_source_extent = SourceExtent(
+            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=3, byteno=2)
+        )
+        self.assertEqual(
+            parse(Peekable(iter([left_paren, int_lit, right_paren]))).source_extent, parenthesized_int_lit_source_extent
         )
 
     def test_parse_spread_preserves_source_extent(self) -> None:
-        ellipsis = Operator("...")
-        ellipsis.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=3, byteno=2)
-        )
-        name = Name("x")
-        name.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=4, byteno=3), end=SourceLocation(lineno=1, colno=4, byteno=3)
+        ellipsis = Operator("...").with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=3, byteno=2)
+            )
         )
-        spread = make_source_annotated_object(
-            Spread,
+        name = Name("x").with_source(
             SourceExtent(
-                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=4, byteno=3)
-            ),
-            "x",
+                start=SourceLocation(lineno=1, colno=4, byteno=3), end=SourceLocation(lineno=1, colno=4, byteno=3)
+            )
         )
-        self.assertTrue(parse(Peekable(iter([ellipsis, name]))).source_extent == spread.source_extent)
+        spread_source_extent = SourceExtent(
+            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=4, byteno=3)
+        )
+        self.assertEqual(parse(Peekable(iter([ellipsis, name]))).source_extent, spread_source_extent)
 
     def test_parse_binop_preserves_source_extent(self) -> None:
-        first_addend = IntLit(1)
-        first_addend.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
-        )
-        operator = Operator("+")
-        operator.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+        first_addend = IntLit(1).with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
+            )
         )
-        second_addend = IntLit(2)
-        second_addend.source_extent = SourceExtent(
-            start=SourceLocation(lineno=2, colno=5, byteno=4), end=SourceLocation(lineno=2, colno=5, byteno=4)
+        operator = Operator("+").with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+            )
         )
-        first_addend_ast = make_source_annotated_object(Int, first_addend.source_extent, 1)
-        operator_ast = BinopKind.ADD
-        second_addend_ast = make_source_annotated_object(Int, second_addend.source_extent, 2)
-        binop = make_source_annotated_object(
-            Binop,
+        second_addend = IntLit(2).with_source(
             SourceExtent(
-                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=2, colno=5, byteno=4)
-            ),
-            operator_ast,
-            first_addend_ast,
-            second_addend_ast,
+                start=SourceLocation(lineno=2, colno=5, byteno=4), end=SourceLocation(lineno=2, colno=5, byteno=4)
+            )
         )
-        self.assertTrue(
-            parse(Peekable(iter([first_addend, operator, second_addend]))).source_extent == binop.source_extent
+        binop_source_extent = SourceExtent(
+            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=2, colno=5, byteno=4)
+        )
+        self.assertEqual(
+            parse(Peekable(iter([first_addend, operator, second_addend]))).source_extent, binop_source_extent
         )
 
     def test_parse_list_preserves_source_extent(self) -> None:
-        left_bracket = LeftBracket()
-        left_bracket.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
-        )
-        one = IntLit(1)
-        one.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+        left_bracket = LeftBracket().with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=1, byteno=0)
+            )
         )
-        comma = Operator(",")
-        comma.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+        one = IntLit(1).with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=2, byteno=1), end=SourceLocation(lineno=1, colno=2, byteno=1)
+            )
         )
-        two = IntLit(2)
-        two.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=5, byteno=4), end=SourceLocation(lineno=1, colno=5, byteno=4)
+        comma = Operator(",").with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=3, byteno=2), end=SourceLocation(lineno=1, colno=3, byteno=2)
+            )
         )
-        right_bracket = RightBracket()
-        right_bracket.source_extent = SourceExtent(
-            start=SourceLocation(lineno=1, colno=6, byteno=5), end=SourceLocation(lineno=1, colno=6, byteno=5)
+        two = IntLit(2).with_source(
+            SourceExtent(
+                start=SourceLocation(lineno=1, colno=5, byteno=4), end=SourceLocation(lineno=1, colno=5, byteno=4)
+            )
         )
-        one_ast = make_source_annotated_object(Int, one.source_extent, 1)
-        two_ast = make_source_annotated_object(Int, two.source_extent, 2)
-        list_ast = make_source_annotated_object(
-            List,
+        right_bracket = RightBracket().with_source(
             SourceExtent(
-                start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=6, byteno=5)
-            ),
-            [one_ast, two_ast],
+                start=SourceLocation(lineno=1, colno=6, byteno=5), end=SourceLocation(lineno=1, colno=6, byteno=5)
+            )
         )
-        self.assertTrue(
-            parse(Peekable(iter([left_bracket, one, comma, two, right_bracket]))).source_extent
-            == list_ast.source_extent
+        list_source_extent = SourceExtent(
+            start=SourceLocation(lineno=1, colno=1, byteno=0), end=SourceLocation(lineno=1, colno=6, byteno=5)
+        )
+        self.assertEqual(
+            parse(Peekable(iter([left_bracket, one, comma, two, right_bracket]))).source_extent, list_source_extent
         )
 
 
@@ -2620,7 +2596,7 @@ def test_eval_count_bits_function_preserves_source_extents(self) -> None:
             match_function_three_source_extent,
         ]
 
-        self.assertTrue(outer_function.source_extent == outer_function_source_extent)
+        self.assertEqual(outer_function.source_extent, outer_function_source_extent)
         self.assertTrue(
             all(
                 match_function.source_extent == source_extent
@@ -2672,9 +2648,9 @@ def test_eval_collatz_function_preserves_source_extents(self) -> None:
             start=SourceLocation(lineno=4, colno=34, byteno=84), end=SourceLocation(lineno=5, colno=79, byteno=205)
         )
 
-        self.assertTrue(outer_function.source_extent == outer_function_source_extent)
-        self.assertTrue(arg.source_extent == arg_source_extent)
-        self.assertTrue(func.source_extent == func_source_extent)
+        self.assertEqual(outer_function.source_extent, outer_function_source_extent)
+        self.assertEqual(arg.source_extent, arg_source_extent)
+        self.assertEqual(func.source_extent, func_source_extent)
 
 
 class ClosureOptimizeTests(unittest.TestCase):