diff --git a/cython/gtsam/tests/test_logging_optimizer.py b/cython/gtsam/tests/test_logging_optimizer.py index c857a6f7d..69665db65 100644 --- a/cython/gtsam/tests/test_logging_optimizer.py +++ b/cython/gtsam/tests/test_logging_optimizer.py @@ -4,6 +4,12 @@ Author: Jing Wu and Frank Dellaert """ # pylint: disable=invalid-name +import sys +if sys.version_info.major >= 3: + from io import StringIO +else: + from cStringIO import StringIO + import unittest from datetime import datetime @@ -37,11 +43,19 @@ class TestOptimizeComet(GtsamTestCase): self.optimizer = gtsam.GaussNewtonOptimizer( graph, initial, self.params) + # setup output capture + self.capturedOutput = StringIO() + sys.stdout = self.capturedOutput + + def tearDown(self): + """Reset print capture.""" + sys.stdout = sys.__stdout__ + def test_simple_printing(self): """Test with a simple hook.""" # Provide a hook that just prints - def hook(_, error: float): + def hook(_, error): print(error) # Only thing we require from optimizer is an iterate method @@ -65,7 +79,7 @@ class TestOptimizeComet(GtsamTestCase): + str(time.hour)+":"+str(time.minute)+":"+str(time.second)) # I want to do some comet thing here - def hook(optimizer, error: float): + def hook(optimizer, error): comet.log_metric("Karcher error", error, optimizer.iterations()) @@ -76,4 +90,4 @@ class TestOptimizeComet(GtsamTestCase): self.gtsamAssertEquals(actual.atRot3(KEY), self.expected) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/cython/gtsam/utils/logging_optimizer.py b/cython/gtsam/utils/logging_optimizer.py index b201bb8aa..939453927 100644 --- a/cython/gtsam/utils/logging_optimizer.py +++ b/cython/gtsam/utils/logging_optimizer.py @@ -4,15 +4,11 @@ Author: Jing Wu and Frank Dellaert """ # pylint: disable=invalid-name -from typing import TypeVar - from gtsam import NonlinearOptimizer, NonlinearOptimizerParams import gtsam -T = TypeVar('T') - -def optimize(optimizer: T, check_convergence, hook): +def optimize(optimizer, check_convergence, hook): """ Given an optimizer and a convergence check, iterate until convergence. After each iteration, hook(optimizer, error) is called. After the function, use values and errors to get the result. @@ -36,8 +32,8 @@ def optimize(optimizer: T, check_convergence, hook): current_error = new_error -def gtsam_optimize(optimizer: NonlinearOptimizer, - params: NonlinearOptimizerParams, +def gtsam_optimize(optimizer, + params, hook): """ Given an optimizer and params, iterate until convergence. After each iteration, hook(optimizer) is called.