Switch take to get

Author: saulshanabrook

python/egg_smol/bindings.pyi

@@ -11,8 +11,8 @@ class EGraph:
     def __init__(self, fact_directory: str | Path | None = None, seminaive=True) -> None: ...
     def parse_program(self, __input: str, /) -> list[_Command]: ...
     def run_program(self, *commands: _Command) -> list[str]: ...
-    def take_extract_report(self) -> Optional[ExtractReport]: ...
-    def take_run_report(self) -> Optional[RunReport]: ...
+    def extract_report(self) -> Optional[ExtractReport]: ...
+    def run_report(self) -> Optional[RunReport]: ...
 
 @final
 class EggSmolError(Exception):

python/egg_smol/egraph.py

@@ -126,7 +126,7 @@ def _simplify(self, expr: EXPR, limit: int, ruleset: Optional[Ruleset], until: t
         tp, decl = expr_parts(expr)
         egg_expr = decl.to_egg(self._decls)
         self._run_program([bindings.Simplify(egg_expr, Config(limit, ruleset, until)._to_egg_config(self._decls))])
-        extract_report = self._get_egraph().take_extract_report()
+        extract_report = self._get_egraph().extract_report()
         if not extract_report:
             raise ValueError("No extract report saved")
         new_tp, new_decl = tp_and_expr_decl_from_egg(self._decls, extract_report.expr)
@@ -181,7 +181,7 @@ def run(self, limit_or_schedule: int | Schedule, /, *until: Fact) -> bindings.Ru
 
     def _run_schedule(self, schedule: Schedule) -> bindings.RunReport:
         self._run_program([bindings.RunScheduleCommand(schedule._to_egg(self._decls))])
-        run_report = self._get_egraph().take_run_report()
+        run_report = self._get_egraph().run_report()
         if not run_report:
             raise ValueError("No run report saved")
         return run_report
@@ -226,7 +226,7 @@ def extract_multiple(self, expr: EXPR, n: int) -> list[EXPR]:
 
     def _run_extract(self, expr: bindings._Expr, n: int) -> bindings.ExtractReport:
         self._run_program([bindings.Extract(n, expr)])
-        extract_report = self._get_egraph().take_extract_report()
+        extract_report = self._get_egraph().extract_report()
         if not extract_report:
             raise ValueError("No extract report saved")
         return extract_report

python/tests/test_bindings.py

@@ -111,7 +111,7 @@ def test_run_rules(self):
         )
         end_time = datetime.datetime.now()
 
-        run_report = egraph.take_run_report()
+        run_report = egraph.run_report()
         assert isinstance(run_report, RunReport)
         total_time = run_report.search_time + run_report.apply_time + run_report.rebuild_time
         # Verify  less than the total time (which includes time spent in Python).
@@ -126,15 +126,15 @@ def test_extract(self):
             Define("y", Call("Num", [Lit(Int(2))]), 1),
             Extract(0, Var("x")),
         )
-        assert egraph.take_extract_report() == ExtractReport(6, Call("Num", [Lit(Int(1))]), [])
+        assert egraph.extract_report() == ExtractReport(6, Call("Num", [Lit(Int(1))]), [])
         egraph.run_program(Extract(0, Var("y")))
         pytest.xfail(reason="https://github.com/mwillsey/egg-smol/issues/128")
-        assert egraph.take_extract_report() == ExtractReport(1, Call("y", []), [])
+        assert egraph.extract_report() == ExtractReport(1, Call("y", []), [])
 
     def test_extract_string(self):
         egraph = EGraph()
         egraph.run_program(Define("x", Lit(String("hello")), None), Extract(0, Var("x")))
-        assert egraph.take_extract_report() == ExtractReport(0, Lit(String("hello")), [])
+        assert egraph.extract_report() == ExtractReport(0, Lit(String("hello")), [])
 
     def test_sort_alias(self):
         # From map example
@@ -149,7 +149,7 @@ def test_sort_alias(self):
             Check([Eq([Lit(String("one")), Call("get", [Var("my_map1"), Lit(Int(1))])])]),
             Extract(0, Var("my_map2")),
         )
-        assert egraph.take_extract_report() == ExtractReport(
+        assert egraph.extract_report() == ExtractReport(
             0,
             Call(
                 "insert",

src/egraph.rs

@@ -50,19 +50,25 @@ impl EGraph {
         Ok(res)
     }
 
-    /// Takes the last expressions extracted from the EGraph, if the last command
+    /// Gets the last expressions extracted from the EGraph, if the last command
     /// was a Simplify or Extract command.
     #[pyo3(signature = ())]
-    fn take_extract_report(&mut self) -> Option<ExtractReport> {
-        info!("Taking last extract report");
-        self.egraph.take_extract_report().map(|r| r.into())
+    fn extract_report(&mut self) -> Option<ExtractReport> {
+        info!("Getting last extract report");
+        match self.egraph.get_extract_report() {
+            Some(report) => Some(report.into()),
+            None => None,
+        }
     }
 
-    /// Takes the last run report from the EGraph, if the last command
+    /// Gets the last run report from the EGraph, if the last command
     /// was a run or simplify command.
     #[pyo3(signature = ())]
-    fn take_run_report(&mut self) -> Option<RunReport> {
-        info!("Taking last run report");
-        self.egraph.take_run_report().map(|r| r.into())
+    fn run_report(&mut self) -> Option<RunReport> {
+        info!("Getting last run report");
+        match self.egraph.get_run_report() {
+            Some(report) => Some(report.into()),
+            None => None,
+        }
     }
 }