Compute win32 python2 length according to encoded size

asottile · asottile · commit ead906aed066 · 2018-10-23T22:02:48.000-07:00
diff --git a/pre_commit/xargs.py b/pre_commit/xargs.py
@@ -3,6 +3,8 @@
 
 import sys
 
+import six
+
 from pre_commit import parse_shebang
 from pre_commit.util import cmd_output
 
@@ -19,9 +21,13 @@ def _command_length(*cmd):
     # win32 uses the amount of characters, more details at:
     # https://github.com/pre-commit/pre-commit/pull/839
     if sys.platform == 'win32':
-        return len(full_cmd.encode('utf-16le')) // 2
-
-    return len(full_cmd.encode(sys.getfilesystemencoding()))
+        # the python2.x apis require bytes, we encode as UTF-8
+        if six.PY2:
+            return len(full_cmd.encode('utf-8'))
+        else:
+            return len(full_cmd.encode('utf-16le')) // 2
+    else:
+        return len(full_cmd.encode(sys.getfilesystemencoding()))
 
 
 class ArgumentTooLongError(RuntimeError):
diff --git a/tests/xargs_test.py b/tests/xargs_test.py
@@ -2,26 +2,36 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
 
+import sys
+
 import mock
 import pytest
+import six
 
 from pre_commit import xargs
 
 
 @pytest.fixture
-def sys_win32_mock():
-    return mock.Mock(
-        platform='win32',
-        getfilesystemencoding=mock.Mock(return_value='utf-8'),
-    )
+def win32_py2_mock():
+    with mock.patch.object(sys, 'getfilesystemencoding', return_value='utf-8'):
+        with mock.patch.object(sys, 'platform', 'win32'):
+            with mock.patch.object(six, 'PY2', True):
+                yield
 
 
 @pytest.fixture
-def sys_linux_mock():
-    return mock.Mock(
-        platform='linux',
-        getfilesystemencoding=mock.Mock(return_value='utf-8'),
-    )
+def win32_py3_mock():
+    with mock.patch.object(sys, 'getfilesystemencoding', return_value='utf-8'):
+        with mock.patch.object(sys, 'platform', 'win32'):
+            with mock.patch.object(six, 'PY2', False):
+                yield
+
+
+@pytest.fixture
+def linux_mock():
+    with mock.patch.object(sys, 'getfilesystemencoding', return_value='utf-8'):
+        with mock.patch.object(sys, 'platform', 'linux'):
+            yield
 
 
 def test_partition_trivial():
@@ -53,31 +63,33 @@ def test_partition_limits():
     )
 
 
-def test_partition_limit_win32(sys_win32_mock):
+def test_partition_limit_win32_py3(win32_py3_mock):
     cmd = ('ninechars',)
     # counted as half because of utf-16 encode
     varargs = ('😑' * 5,)
-    with mock.patch('pre_commit.xargs.sys', sys_win32_mock):
-        ret = xargs.partition(cmd, varargs, _max_length=20)
+    ret = xargs.partition(cmd, varargs, _max_length=20)
+    assert ret == (cmd + varargs,)
+
 
+def test_partition_limit_win32_py2(win32_py2_mock):
+    cmd = ('ninechars',)
+    varargs = ('😑' * 5,)  # 4 bytes * 5
+    ret = xargs.partition(cmd, varargs, _max_length=30)
     assert ret == (cmd + varargs,)
 
 
-def test_partition_limit_linux(sys_linux_mock):
+def test_partition_limit_linux(linux_mock):
     cmd = ('ninechars',)
     varargs = ('😑' * 5,)
-    with mock.patch('pre_commit.xargs.sys', sys_linux_mock):
-        ret = xargs.partition(cmd, varargs, _max_length=30)
-
+    ret = xargs.partition(cmd, varargs, _max_length=30)
     assert ret == (cmd + varargs,)
 
 
-def test_argument_too_long_with_large_unicode(sys_linux_mock):
+def test_argument_too_long_with_large_unicode(linux_mock):
     cmd = ('ninechars',)
     varargs = ('😑' * 10,)  # 4 bytes * 10
-    with mock.patch('pre_commit.xargs.sys', sys_linux_mock):
-        with pytest.raises(xargs.ArgumentTooLongError):
-            xargs.partition(cmd, varargs, _max_length=20)
+    with pytest.raises(xargs.ArgumentTooLongError):
+        xargs.partition(cmd, varargs, _max_length=20)
 
 
 def test_argument_too_long():
