Begin specializing support for flipping endianness (#3420)

Author: ESultanik

polyfile/magic.py

@@ -877,6 +877,9 @@ def all_extensions(self) -> LazyIterableSet[str]:
         """Returns the set of all possible extensions that this test or any of its descendants could match against"""
         return LazyIterableSet(self._all_extensions())
 
+    def test_flip_endianness(self, data: bytes, absolute_offset: int, parent_match: Optional[TestResult]) -> TestResult:
+        raise NotImplementedError(f"TODO: Implement test_flip_endianness for {self.__class__.__name__}")
+
     @abstractmethod
     def test(self, data: bytes, absolute_offset: int, parent_match: Optional[TestResult]) -> TestResult:
         raise NotImplementedError()
@@ -925,14 +928,22 @@ def write(self, writer: ANSIWriter, is_current_test: bool = False, pre_mime_text
     def calculate_absolute_offset(self, data: bytes, parent_match: Optional[TestResult] = None) -> int:
         return self.offset.to_absolute(data, parent_match)
 
-    def _match(self, context: MatchContext, parent_match: Optional[TestResult] = None) -> Iterator[MatchedTest]:
+    def _match(
+            self,
+            context: MatchContext,
+            parent_match: Optional[TestResult] = None,
+            flip_endianness: bool = False
+    ) -> Iterator[MatchedTest]:
         if context.only_match_mime and not self.can_match_mime:
             return
         try:
             absolute_offset = self.calculate_absolute_offset(context.data, parent_match)
         except InvalidOffsetError:
             return
-        m = self.test(context.data, absolute_offset, parent_match)
+        if flip_endianness:
+            m = self.test_flip_endianness(context.data, absolute_offset, parent_match)
+        else:
+            m = self.test(context.data, absolute_offset, parent_match)
         if logging.root.level <= TRACE and (bool(m) or self.level > 0):
             log.trace(
                 f"{self.source_info!s}\t{bool(m)}\t{absolute_offset}\t"
@@ -943,7 +954,7 @@ def _match(self, context: MatchContext, parent_match: Optional[TestResult] = Non
                 yield m
             for child in self.children:
                 if not context.only_match_mime or child.can_match_mime:
-                    yield from child._match(context=context, parent_match=m)
+                    yield from child._match(context=context, parent_match=m, flip_endianness=flip_endianness)
 
     def match(self, to_match: Union[bytes, BinaryIO, str, Path, MatchContext]) -> Iterator[TestResult]:
         """Yields all matches for the given data"""
@@ -2235,6 +2246,13 @@ def to_absolute(self, data: bytes, last_match: Optional[TestResult], allow_inval
     def subtest_type(self) -> TestType:
         return TestType.UNKNOWN
 
+    def test_flip_endianness(self, data: bytes, absolute_offset: int, parent_match: Optional[TestResult]) -> TestResult:
+        if parent_match is not None:
+            return MatchedTest(self, offset=parent_match.offset + parent_match.length, length=0, value=self.name,
+                               parent=parent_match)
+        else:
+            raise ValueError("A named test must always be called from a `use` test.")
+
     def test(self, data: bytes, absolute_offset: int, parent_match: Optional[TestResult]) -> MatchedTest:
         if parent_match is not None:
             return MatchedTest(self, offset=parent_match.offset + parent_match.length, length=0, value=self.name,
@@ -2273,9 +2291,13 @@ def referenced_tests(self) -> Set[NamedTest]:
             result |= self.referenced_test.referenced_tests()
         return result
 
-    def _match(self, context: MatchContext, parent_match: Optional[TestResult] = None) -> Iterator[TestResult]:
-        if self.flip_endianness:
-            raise NotImplementedError("TODO: Add support for use tests with flipped endianness")
+    def _match(
+            self,
+            context: MatchContext,
+            parent_match: Optional[TestResult] = None,
+            flip_endianness: bool = False
+    ) -> Iterator[TestResult]:
+        flip_endianness = flip_endianness ^ self.flip_endianness
         try:
             absolute_offset = self.offset.to_absolute(context.data, last_match=parent_match)
         except InvalidOffsetError:
@@ -2285,7 +2307,7 @@ def _match(self, context: MatchContext, parent_match: Optional[TestResult] = Non
         )
         use_match = MatchedTest(self, None, absolute_offset, 0, parent=parent_match)
         yielded = False
-        for named_result in self.referenced_test._match(context, use_match):
+        for named_result in self.referenced_test._match(context, use_match, flip_endianness=flip_endianness):
             if not yielded:
                 yielded = True
                 yield use_match
@@ -2298,7 +2320,7 @@ def _match(self, context: MatchContext, parent_match: Optional[TestResult] = Non
             return
         for child in self.children:
             if not context.only_match_mime or child.can_match_mime:
-                yield from child._match(context=context, parent_match=use_match)
+                yield from child._match(context=context, parent_match=use_match, flip_endianness=flip_endianness)
 
     def test(self, data: bytes, absolute_offset: int, parent_match: Optional[TestResult]) -> TestResult:
         raise NotImplementedError("This function should never be called")