Added os_version, Python 3 fixes

3 files changed, 51 insertions, 2 deletions

diff --git a/machammer/functions.py b/machammer/functions.py
index fa4fc63..326918d 100644
--- a/machammer/functions.py
+++ b/machammer/functions.py
@@ -293,3 +293,44 @@ def create_os_media(src, dst):
 
 def log(msg):
     logging.debug(msg)
+
+
+def os_version():
+    class OsVersion(object):
+        def __init__(self):
+            out = {}
+            for _ in check_output('/usr/bin/sw_vers').decode().split("\n"):
+                (k, v) = _.replace(':', '').split("\t")
+                out[k] = v
+
+            self.name = out['ProductName']
+            self.build = out['BuildVersion']
+            self.version = out['ProductVersion']
+            self._version = self.pad(self.version)
+
+        def pad(self, v):
+            return int(str(v).replace('.', '').ljust(8, '0'))
+
+        def __repr__(self):
+            return '%s (%s)' % (self.version, self.build)
+
+        def __eq__(self, v):
+            return self.version == v
+
+        def __ne__(self, v):
+            return self.version != v
+
+        def __lt__(self, v):
+            return self._version < self.pad(v)
+
+        def __gt__(self, v):
+            print(self._version, self.pad(v))
+            return self._version > self.pad(v)
+
+        def __ge__(self, v):
+            return self._version >= self.pad(v)
+
+        def __str__(self):
+            return self.version
+
+    return OsVersion()
diff --git a/machammer/system_profiler.py b/machammer/system_profiler.py
index 81dfb69..711ba24 100644
--- a/machammer/system_profiler.py
+++ b/machammer/system_profiler.py
@@ -20,7 +20,9 @@ PROFILER_PATH = '/usr/sbin/system_profiler'
 class SystemProfile(object):
     def __init__(self, dt=DEFAULT_DT):
         types = subprocess.check_output([PROFILER_PATH, '-listDataTypes']).strip()
-        self.types = [x[2:].replace('DataType', '') for x in types.split("\n") if x.startswith('SP')]
+        types = types.decode().split("\n")
+
+        self.types = [x[2:].replace('DataType', '') for x in types if x.startswith('SP')]
         self.types.sort()
 
         if dt not in self.types:
diff --git a/tests.py b/tests.py
index 95288b7..c85ac26 100755
--- a/tests.py
+++ b/tests.py
@@ -137,12 +137,18 @@ class AppsTestCase(TestCase):
 
     def testStickiesVersion(self):
         results = self.profile.find('_name', 'Stickies')
-        self.assertEquals(results[0]['version'], '10.0')
+        if functions.os_version() >= 10.13:
+            self.assertEqual(results[0]['version'], '10.1')
+        else:
+            self.assertEqual(results[0]['version'], '10.0')
 
     def testFindApplications(self):
         results = self.profile.find('path', '/Applications')
         self.assertTrue(len(results) > 10)
 
+    def testSystemVersion(self):
+        self.assertLess(functions.os_version(), 10.14)
+
 
 class InstallerTestCase(TestCase):
     def setUp(self):