Merge pull request #555 from ChrisCummins/fix/benchmark-uri

[datasets] Small benchmark URI fix and tweak

compiler_gym/datasets/uri.py

@@ -130,6 +130,12 @@ def from_string(cls, uri: str) -> "BenchmarkUri":
             fragment=components.fragment,
         )
 
+    def startswith(self, *args):
+        return str(self).startswith(*args)
+
+    def endswith(self, *args):
+        return str(self).endswith(*args)
+
     def __repr__(self):
         return urlunparse(
             ParseResult(

examples/example_compiler_gym_service/__init__.py

@@ -72,7 +72,7 @@ def __init__(self, *args, **kwargs):
     def benchmark_uris(self) -> Iterable[str]:
         yield from (f"benchmark://example-v0{k}" for k in self._benchmarks.keys())
 
-    def benchmark_from_parsed_uris(self, uri: BenchmarkUri) -> Benchmark:
+    def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
         if uri.path in self._benchmarks:
             return self._benchmarks[uri]
         else:

tests/datasets/uri_test.py

@@ -95,5 +95,19 @@ def test_canonicalize_1():
     assert BenchmarkUri.canonicalize("test-v0") == "benchmark://test-v0"
 
 
+def test_startswith():
+    uri = BenchmarkUri.from_string("benchmark://test-v0/foo")
+    assert not uri.startswith("!!!")
+    assert uri.startswith("b")
+    assert uri.startswith("benchmark://test-v0/fo")
+
+
+def test_endswith():
+    uri = BenchmarkUri.from_string("benchmark://test-v0/foo")
+    assert not uri.endswith("!!!")
+    assert uri.endswith("o")
+    assert uri.endswith("mark://test-v0/foo")
+
+
 if __name__ == "__main__":
     main()