diff --git a/geist/backends/fake.py b/geist/backends/fake.py index 3f37fba..c357519 100644 --- a/geist/backends/fake.py +++ b/geist/backends/fake.py @@ -1,4 +1,6 @@ import numpy as np +from six import string_types + from ..finders import Location, LocationList @@ -14,7 +16,7 @@ def __init__(self, **kwargs): [Location(0, 0, w=w, h=h, image=self.image)] ) else: - if isinstance(image, basestring): + if isinstance(image, string_types): image = np.load(image) self.image = image h, w, _ = image.shape diff --git a/geist/backends/replay.py b/geist/backends/replay.py index 6d98b1a..78c1392 100644 --- a/geist/backends/replay.py +++ b/geist/backends/replay.py @@ -3,7 +3,7 @@ import sys import json import base64 -import StringIO +from six import StringIO import wrapt from PIL import Image from . import get_platform_backend @@ -119,7 +119,7 @@ def capture_locations(self): location_list = LocationList() for b64_location in b64_locations: base64_png = b64_location['base64_png'] - string_file = StringIO.StringIO( + string_file = StringIO( base64.b64decode(base64_png) ) x, y, w, h = ( @@ -191,7 +191,7 @@ def _write_action(self, funcname, *args, **kwargs): def _write_capture_locations(self, locations): b64_locations = [] for location in locations: - string_file = StringIO.StringIO() + string_file = StringIO() Image.fromarray(location.image).save(string_file, 'png') b64_png = base64.b64encode(string_file.getvalue()) b64_locations.append({ diff --git a/geist/backends/xvfb.py b/geist/backends/xvfb.py index 4390813..3bbcbd9 100644 --- a/geist/backends/xvfb.py +++ b/geist/backends/xvfb.py @@ -57,7 +57,7 @@ def __init__(self, **kwargs): display = ":%d" % (self.display_num, ) self._display_dir = XVFB_PATH % (self.display_num,) - os.makedirs(self._display_dir, 0700) + os.makedirs(self._display_dir, 0o700) dev_null = open('/dev/null', 'w') self._xvfb_proc = subprocess.Popen( [ diff --git a/geist/pyplot.py b/geist/pyplot.py index 6ba2db3..d488b29 100644 --- a/geist/pyplot.py +++ b/geist/pyplot.py @@ -68,7 +68,6 @@ def _get_colour(self, numpy_array): smax = sat.max() vmin = val.min() vmax = val.max() - print hmin, hmax, smin, smax, vmin, vmax return hsv(lambda h, s, v: ( (h >= hmin) & (h <= hmax)) & ((s >= smin) & (s <= smax)) & diff --git a/geist/vision.py b/geist/vision.py index d8984bd..1be92af 100644 --- a/geist/vision.py +++ b/geist/vision.py @@ -1,7 +1,8 @@ -from __future__ import division +from __future__ import print_function, division from numpy.fft import irfft2, rfft2 import numpy import itertools +from six import iteritems import operator @@ -84,10 +85,10 @@ def best_convolution(bin_template, bin_image, max_hor_cells = iw // th # Try to work out how many times we can stack the image - usable_factors = {n: factors for n, factors in overlap_table.iteritems() + usable_factors = {n: factors for n, factors in iteritems(overlap_table) if ((template_sum + 1) ** (n)) < ACCURACY_LIMIT} overlap_options = [(factor, n // factor) - for n, factors in usable_factors.iteritems() + for n, factors in iteritems(usable_factors) for factor in factors if (factor <= max_vert_cells and n // factor <= max_hor_cells)] @@ -526,5 +527,5 @@ def correlation_coefficient_normed(template, image): (template_distance ** 2).sum() * (image_distance_of_template_area ** 2).sum() ) - print y + print(y) return corr_num / corr_denum diff --git a/geist_tests/test_finders.py b/geist_tests/test_finders.py index 952c3ac..b769f29 100644 --- a/geist_tests/test_finders.py +++ b/geist_tests/test_finders.py @@ -18,3 +18,11 @@ def test_filter(self): result = self.gui.wait_find_one(finder) self.assertEqual(parent_x + child_x, result.x) self.assertEqual(parent_y + child_y, result.y) + + +location_finder_filter_suite = unittest.TestLoader().loadTestsFromTestCase( + TestLocationFinderFilter) +all_tests = unittest.TestSuite([location_finder_filter_suite]) +if __name__ == "__main__": + runner = unittest.TextTestRunner(verbosity=1) + runner.run(all_tests) diff --git a/setup.py b/setup.py index 7325605..95eead2 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ 'pillow', 'wrapt', 'mock', + 'six', ], description='Visual Automation Library', classifiers=[