shrinkray 0.0.0__py3-none-any.whl

shrinkray/learning.py

@@ -0,0 +1,221 @@
+import hashlib
+import os
+import random
+import re
+import site
+import subprocess
+import sys
+from glob import glob
+from typing import Awaitable, Callable
+
+import trio
+from attrs import define
+
+from shrinkray.problem import BasicReductionProblem
+from shrinkray.reducer import ShrinkRay
+from shrinkray.work import WorkContext
+
+WHITESPACE = re.compile(rb"\s+")
+
+
+def whitespace_normalize(s):
+    return WHITESPACE.sub(b" ", s).strip()
+
+
+def threshold_value(seed: bytes, test_case: bytes) -> float:
+    return random.Random(hashlib.sha1(seed + b":" + test_case).digest()).random()
+
+
+@define
+class RandomInterestingnessTest:
+    seed: bytes
+    threshold: float
+    key_substring: bytes
+    base_interestingness: Callable[[bytes], Awaitable[bool]]
+
+    async def __call__(self, test_case: bytes) -> bool:
+        if threshold_value(self.seed, test_case) > self.threshold:
+            await trio.lowlevel.checkpoint()
+            return False
+        normalized = whitespace_normalize(test_case)
+        if self.key_substring not in normalized:
+            await trio.lowlevel.checkpoint()
+            return False
+        return await self.base_interestingness(test_case)
+
+
+class CouldNotLearn(Exception):
+    pass
+
+
+class CorpusExplorer:
+    def __init__(
+        self,
+        is_interesting: Callable[[bytes], Awaitable[bool]],
+        corpus: list[str],
+        parallelism=os.cpu_count(),
+    ):
+        self.__is_interesting = is_interesting
+        self.__is_interesting_cache = {}
+        self.__corpus_files = [t for t in corpus if os.stat(t).st_size > 0]
+        self.__corpus_data = []
+        self.__limit = trio.CapacityLimiter(parallelism)
+        self.__wait_for_me = object()
+        self.__random = random.Random()
+        self.__parallelism = parallelism
+
+    async def is_interesting(self, tc):
+        try:
+            result = self.__is_interesting_cache[tc]
+            if result != self.__wait_for_me:
+                await trio.lowlevel.checkpoint()
+                return result
+            while result == self.__wait_for_me:
+                await trio.sleep(0.01)
+                result = self.__is_interesting_cache[tc]
+            assert result is not self.__wait_for_me
+            return result
+        except KeyError:
+            self.__is_interesting_cache[tc] = self.__wait_for_me
+            async with self.__limit:
+                result = await self.__is_interesting(tc)
+            self.__is_interesting_cache[tc] = result
+            return result
+
+    async def random_corpus_member(self):
+        return list(await self.random_corpus_sample(1))[0]
+
+    async def random_corpus_sample(self, max_elements=100):
+        result = []
+        max_from_data = len(self.__corpus_data)
+
+        async def run():
+            nonlocal max_from_data
+            while len(result) < max_elements and (
+                max_from_data > 0 or self.__corpus_files
+            ):
+                i = self.__random.randrange(0, len(self.__corpus_files) + max_from_data)
+                if i < len(self.__corpus_files):
+                    j = len(self.__corpus_files) - 1
+                    self.__corpus_files[i], self.__corpus_files[j] = (
+                        self.__corpus_files[j],
+                        self.__corpus_files[i],
+                    )
+                    f = self.__corpus_files.pop()
+                    with open(f, "rb") as input:
+                        data = input.read()
+
+                    if await self.is_interesting(data):
+                        self.__corpus_data.append(data)
+                        result.append(data)
+                else:
+                    i -= len(self.__corpus_files)
+                    j = max_from_data - 1
+                    self.__corpus_data[i], self.__corpus_data[j] = (
+                        self.__corpus_data[j],
+                        self.__corpus_data[i],
+                    )
+                    result.append(self.__corpus_data[j])
+                    max_from_data -= 1
+                    await trio.lowlevel.checkpoint()
+
+        if self.__parallelism > 1:
+            async with trio.open_nursery() as nursery:
+                for _ in range(self.__parallelism):
+                    nursery.start_soon(run)
+        else:
+            await run()
+        return result[:max_elements]
+
+    async def random_reduction_problem(self):
+        while True:
+            sample = await self.random_corpus_sample(400)
+
+            index = defaultdict(list)
+
+            for i, s in enumerate(sample):
+                for j, c in enumerate(whitespace_normalize(s)):
+                    index[c].append((i, j))
+
+            queue = dequeue()
+
+            strings = [b""]
+
+            if not normalized:
+                continue
+            while True:
+                start = self.__random.randrange(0, len(normalized))
+                if not WHITESPACE.match(normalized[start : start + 1]):
+                    break
+            while True:
+                end = self.__random.randint(start, min(start + 20, len(normalized) - 1))
+                if not WHITESPACE.match(normalized[end : end + 1]):
+                    break
+            substring = normalized[start:end]
+            while True:
+                seed = str(self.__random.randint(0, 10**6)).encode("ascii")
+                value = threshold_value(seed=seed, test_case=initial)
+                if value >= 0.5:
+                    threshold = self.__random.random() * (1.0 - value) + value
+                    break
+
+            return initial, RandomInterestingnessTest(
+                seed=seed,
+                threshold=threshold,
+                key_substring=substring,
+                base_interestingness=self.is_interesting,
+            )
+
+    async def full_shrink(self, initial, condition) -> bytes:
+        problem: BasicReductionProblem[bytes] = BasicReductionProblem(
+            initial=initial,
+            is_interesting=condition,
+            work=WorkContext(parallelism=self.__parallelism, random=self.__random),
+        )
+
+        reducer = ShrinkRay(
+            target=problem,
+        )
+
+        await reducer.run()
+
+        return problem.current_test_case
+
+    async def all_shrinks(self, sample, condition):
+        shrinks = set()
+
+        async def run_and_add(c, condition):
+            shrinks.add(await self.full_shrink(c, condition))
+
+        async with trio.open_nursery() as nursery:
+            for d in sample:
+                nursery.start_soon(run_and_add, d, condition)
+
+        return shrinks
+
+
+def python_files() -> list[str]:
+    return [
+        f
+        for d in site.getsitepackages()
+        for f in glob(os.path.join(d, "**", "*.py"), recursive=True)
+    ]
+
+
+ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+
+IS_PYTHON_SCRIPT = os.path.join(ROOT, "scripts", "ispython.py")
+
+assert os.path.exists(IS_PYTHON_SCRIPT), ROOT
+
+
+async def is_python(data):
+    return (
+        await trio.run_process(
+            [sys.executable, IS_PYTHON_SCRIPT],
+            stdin=data,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        )
+    ).returncode == 0