upgrade to python3

python/datasource.py

@@ -14,6 +14,31 @@ from engineconfig import getConfig
 from entity import Submit, Problem, TestCase, PresetCode, DataFile
 
 
+def _to_bytes(value, encoding='utf-8', errors='replace'):
+    if value is None:
+        return b''
+    if isinstance(value, xmlrpclib.Binary):
+        return value.data
+    if isinstance(value, bytes):
+        return value
+    if isinstance(value, str):
+        return value.encode(encoding, errors)
+    return str(value).encode(encoding, errors)
+
+
+def _to_text(value, encoding='utf-8', errors='replace'):
+    if value is None:
+        return ''
+    if isinstance(value, xmlrpclib.Binary):
+        value = value.data
+    if isinstance(value, bytes):
+        return value.decode(encoding, errors)
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+
 class DataSourceError(Exception):
     pass
 
@@ -244,9 +269,10 @@ class XmlRpcDataSource:
             try:
                 submits = self.server.oj.get_submits(judgeid, limit)
                 for submit in submits:
-                    if not isinstance(submit['code'], str):
-                        submit['code'] = submit['code'].__str__()
+                    if not isinstance(submit.get('code'), str):
+                        submit['code'] = _to_text(submit.get('code'))
                 return submits
+
             except xmlrpclib.Error as e:
                 raise DataSourceError(e)
             except socket.error as e:
@@ -268,12 +294,13 @@ class XmlRpcDataSource:
             try:
                 tests = self.server.oj.get_tests(problemid, full)
                 for test in tests:
-                    if not isinstance(test['input'], str):
-                        test['input'] = test['input'].__str__()
-                    if not isinstance(test['output'], str):
-                        test['output'] = test['output'].__str__()
+                    if not isinstance(test.get('input'), str):
+                        test['input'] = _to_text(test.get('input'))
+                    if not isinstance(test.get('output'), str):
+                        test['output'] = _to_text(test.get('output'))
                 self.logger.debug('Got %d test case(s)', len(tests))
                 return tests
+
             except xmlrpclib.Error as e:
                 raise DataSourceError(e)
             except socket.error as e:
@@ -284,13 +311,16 @@ class XmlRpcDataSource:
         while True:
             try:
                 test = self.server.oj.get_gztest(testid)
-                if not isinstance(test['input'], str):
-                    test['input'] = test['input'].__str__()
-                if not isinstance(test['output'], str):
-                    test['output'] = test['output'].__str__()
-                test['input'] = bz2.decompress(test['input'])
-                test['output'] = bz2.decompress(test['output'])
+                try:
+                    input_data = bz2.decompress(_to_bytes(test.get('input')))
+                    output_data = bz2.decompress(_to_bytes(test.get('output')))
+                except OSError:
+                    self.logger.exception('Failed to decompress test data for %s', testid)
+                    raise DataSourceError('Invalid bz2 data for test %s' % testid)
+                test['input'] = _to_text(input_data)
+                test['output'] = _to_text(output_data)
                 return test
+
             except xmlrpclib.Error as e:
                 raise DataSourceError(e)
             except socket.error as e:
@@ -302,10 +332,11 @@ class XmlRpcDataSource:
             try:
                 codes = self.server.oj.get_presetcodes(problemid, lang)
                 for code in codes:
-                    if not isinstance(code['code'], str):
-                        code['code'] = code['code'].__str__()
+                    if not isinstance(code.get('code'), str):
+                        code['code'] = _to_text(code.get('code'))
                 self.logger.debug('Got %d presetcodes', len(codes))
                 return codes
+
             except xmlrpclib.Error as e:
                 raise DataSourceError(e)
             except socket.error as e:
@@ -328,7 +359,8 @@ class XmlRpcDataSource:
         while True:
             try:
                 data = self.server.oj.get_datafile_data(datafileid)
-                return str(data)
+                return _to_bytes(data)
+
             except xmlrpclib.Error as e:
                 raise DataSourceError(e)
             except socket.error as e:
@@ -336,7 +368,8 @@ class XmlRpcDataSource:
                 time.sleep(self.config.retry_wait)
 
     def update_submit_compilemessage(self, id, compilemsg):
-        compilemsg = xmlrpclib.Binary(compilemsg)
+        compilemsg = xmlrpclib.Binary(_to_bytes(compilemsg))
+
         while True:
             try:
                 return self.server.oj.update_submit_compilemessage(
@@ -349,10 +382,11 @@ class XmlRpcDataSource:
 
     def update_submit_test_results(self, id, results):
         for r in results:
-            if not isinstance(r['stdout'], str): r['stdout'] = ''
-            if not isinstance(r['stderr'], str): r['stderr'] = ''
-            r['stdout'] = xmlrpclib.Binary(r['stdout'])
-            r['stderr'] = xmlrpclib.Binary(r['stderr'])
+            stdout = r.get('stdout', '')
+            stderr = r.get('stderr', '')
+            r['stdout'] = xmlrpclib.Binary(_to_bytes(stdout))
+            r['stderr'] = xmlrpclib.Binary(_to_bytes(stderr))
+
         while True:
             try:
                 return self.server.oj.update_submit_test_results(id, results)
@@ -389,10 +423,8 @@ class XmlRpcDataSource:
         while True:
             try:
                 msg = self.server.oj.get_submit_compilemessage(sid)
-                if isinstance(msg, str):
-                    return msg
-                else:
-                    return msg.__str__()
+                return _to_text(msg)
+
             except xmlrpclib.Error as e:
                 return DataSourceError(e)
             except socket.error as e:
@@ -411,7 +443,16 @@ class XmlRpcDataSource:
 
 class DataSourceTest(unittest.TestCase):
 
+    def testByteTextConversion(self):
+        raw = b'hello'
+        self.assertEqual(_to_text(raw), 'hello')
+        self.assertEqual(_to_bytes('hello'), b'hello')
+        binval = xmlrpclib.Binary(b'world')
+        self.assertEqual(_to_text(binval), 'world')
+        self.assertEqual(_to_bytes(binval), b'world')
+
     def setUp(self):
+
         exec(open(os.path.join('..', 'testdata', 'test_config.py')).read())
         self.config = getConfig()
         self.datasource = self.config.datasources[0]