[cmpcodesize][2] Begin extracting regular expressions

utils/cmpcodesize/cmpcodesize/compare.py

@@ -1,9 +1,10 @@
 import re
 import os
-import subprocess
 import collections
 from operator import itemgetter
 
+from cmpcodesize import otool, regex
+
 Prefixes = {
     # Cpp
     "__Z" : "CPP",
@@ -64,44 +65,27 @@ def addFunction(sizes, function, startAddr, endAddr, groupByPrefix):
         sizes[function] += size
 
 
-def flatten(*args):
-    for x in args:
-        if hasattr(x, '__iter__'):
-            for y in flatten(*x):
-                yield y
-        else:
-            yield x
-
-
 def readSizes(sizes, fileName, functionDetails, groupByPrefix):
     # Check if multiple architectures are supported by the object file.
     # Prefer arm64 if available.
-    architectures = subprocess.check_output(["otool", "-V", "-f", fileName]).split("\n")
-    arch = None
-    archPattern = re.compile('architecture ([\S]+)')
-    for architecture in architectures:
-        archMatch = archPattern.match(architecture)
-        if archMatch:
-            if arch is None:
-                arch = archMatch.group(1)
-            if "arm64" in arch:
-                arch = "arm64"
-    if arch is not None:
-      archParams = ["-arch", arch]
-    else:
-      archParams = []
-
+    fat_headers = otool.fat_headers(fileName)
+    architecture = regex.architecture(fat_headers)
     if functionDetails:
-        content = subprocess.check_output(flatten(["otool", archParams, "-l", "-v", "-t", fileName])).split("\n")
-        content += subprocess.check_output(flatten(["otool", archParams, "-v", "-s", "__TEXT", "__textcoal_nt", fileName])).split("\n")
+        content = otool.load_commands(fileName,
+                                      architecture=architecture,
+                                      include_text_sections=True).split('\n')
+        content += otool.text_sections(fileName,
+                                       architecture=architecture).split('\n')
     else:
-        content = subprocess.check_output(flatten(["otool", archParams, "-l", fileName])).split("\n")
+        content = otool.load_commands(fileName,
+                                      architecture=architecture).split('\n')
 
     sectName = None
     currFunc = None
     startAddr = None
     endAddr = None
 
+    # FIXME: Move re calls into cmpcodesize.regex module.
     sectionPattern = re.compile(' +sectname ([\S]+)')
     sizePattern = re.compile(' +size ([\da-fx]+)')
     asmlinePattern = re.compile('^([0-9a-fA-F]+)\s')

utils/cmpcodesize/cmpcodesize/otool.py

@@ -0,0 +1,45 @@
+import subprocess
+
+
+def _command_for_architecture(architecture=None):
+    if architecture is None:
+        return ['otool']
+    else:
+        return ['otool', '-arch', architecture]
+
+
+def fat_headers(path):
+    """
+    Returns the headers for the executable at the given path.
+    Raises a subprocess.CalledProcessError if otool encounters an error,
+    such as not finding a file at the given path.
+    """
+    return subprocess.check_output(['otool', '-V', '-f', path])
+
+
+def load_commands(path, architecture=None, include_text_sections=False):
+    """
+    Returns the load commands for the executable at the given path,
+    for the given architecture. If print_text_section is specified,
+    the disassembled text section of the load commands is also output.
+
+    Raises a subprocess.CalledProcessError if otool encounters an error,
+    such as not finding a file at the given path.
+    """
+    command = _command_for_architecture(architecture) + ['-l']
+    if include_text_sections:
+        command += ['-v', '-t']
+    return subprocess.check_output(command + [path])
+
+
+def text_sections(path, architecture=None):
+    """
+    Returns the contents of the text sections of the executable at the
+    given path, for the given architecture.
+
+    Raises a subprocess.CalledProcessError if otool encounters an error,
+    such as not finding a file at the given path.
+    """
+    return subprocess.check_output(
+        _command_for_architecture(architecture) +
+        ['-v', '-s', '__TEXT', '__textcoal_nt', path])

utils/cmpcodesize/cmpcodesize/regex.py

@@ -0,0 +1,26 @@
+import re
+
+
+# Cache the compiled regex into a global object.
+_ARCHITECTURE_REGEX = re.compile('architecture (\S+)')
+
+
+def architecture(fat_headers):
+    """
+    Given a string representing fat headers from an executable,
+    returns one of the following:
+
+    1. arm64, if that is one of the architectures listed.
+    2. If arm64 us not listed, the first architecture that is listed.
+    3. None, if no architectures are listed.
+    """
+    result = None
+    for line in fat_headers.splitlines():
+        match = _ARCHITECTURE_REGEX.match(line)
+        if match:
+            arch = match.group(1)
+            if arch == 'arm64':
+                return arch
+            elif result is None:
+                result = match.group(1)
+    return result

utils/cmpcodesize/tests/test_otool.py

@@ -0,0 +1,73 @@
+import subprocess
+import unittest
+
+from cmpcodesize import otool
+
+
+# Store parameters passed to subprocess.check_output into
+# this global variable.
+_subprocess_check_output_arguments = []
+
+
+# We'll monkey-patch subprocess.check_output with this stub
+# function, which simply records whatever's passed to
+# check_output into the global variables above.
+def _stub_subprocess_check_output(arguments, *args, **kwargs):
+    global _subprocess_check_output_arguments
+    _subprocess_check_output_arguments = arguments
+
+
+class OtoolTestCase(unittest.TestCase):
+    def setUp(self):
+        # Monkey-patch subprocess.check_output with our stub function.
+        self._original_check_output = subprocess.check_output
+        subprocess.check_output = _stub_subprocess_check_output
+
+    def tearDown(self):
+        # Undo the monkey-patching.
+        subprocess.check_output = self._original_check_output
+
+    def test_fat_headers(self):
+        otool.fat_headers('/path/to/foo')
+        self.assertEqual(_subprocess_check_output_arguments,
+                         ['otool', '-V', '-f', '/path/to/foo'])
+
+    def test_load_commands_with_no_architecture(self):
+        otool.load_commands('/path/to/bar')
+        self.assertEqual(_subprocess_check_output_arguments,
+                         ['otool', '-l', '/path/to/bar'])
+
+    def test_load_commands_with_architecture(self):
+        otool.load_commands('/path/to/baz', architecture='arch-foo')
+        self.assertEqual(
+            _subprocess_check_output_arguments,
+            ['otool', '-arch', 'arch-foo', '-l', '/path/to/baz'])
+
+    def test_load_commands_no_architecture_but_including_text_sections(self):
+        otool.load_commands(
+            '/path/to/flim', include_text_sections=True)
+        self.assertEqual(
+            _subprocess_check_output_arguments,
+            ['otool', '-l', '-v', '-t', '/path/to/flim'])
+
+    def test_load_commands_with_architecture_and_including_text_sections(self):
+        otool.load_commands(
+            '/path/to/flam',
+            architecture='arch-bar',
+            include_text_sections=True)
+        self.assertEqual(
+            _subprocess_check_output_arguments,
+            ['otool', '-arch', 'arch-bar', '-l', '-v', '-t', '/path/to/flam'])
+
+    def test_text_sections_no_architecture(self):
+        otool.text_sections('/path/to/fish')
+        self.assertEqual(
+            _subprocess_check_output_arguments,
+            ['otool', '-v', '-s', '__TEXT', '__textcoal_nt', '/path/to/fish'])
+
+    def test_text_sections_with_architecture(self):
+        otool.text_sections('/path/to/frosh', architecture='arch-baz')
+        self.assertEqual(
+            _subprocess_check_output_arguments,
+            ['otool', '-arch', 'arch-baz', '-v', '-s',
+             '__TEXT', '__textcoal_nt', '/path/to/frosh'])

utils/cmpcodesize/tests/test_regex.py

@@ -0,0 +1,41 @@
+import unittest
+
+from cmpcodesize import regex
+
+
+class ArchitectureTestCase(unittest.TestCase):
+    def test_no_architectures_listed_returns_none(self):
+        headers = 'Fat headers\n' + \
+                  '    cputype CPU_TYPE_X86_64\n'
+        self.assertIsNone(regex.architecture(headers))
+
+    def test_arm64_listed_returns_arm64(self):
+        headers = 'Fat headers\n' + \
+                  'architecture x86_64\n' + \
+                  '    cputype CPU_TYPE_X86_64\n' + \
+                  'architecture arm64\n' + \
+                  '    cputype CPU_TYPE_ARM64\n'
+        self.assertEquals(regex.architecture(headers), 'arm64')
+
+    def test_arm64_not_listed_returns_first(self):
+        headers = 'Fat headers\n' + \
+                  'architecture x86_64\n' + \
+                  '    cputype CPU_TYPE_X86_64\n' + \
+                  'architecture i386\n' + \
+                  '    cputype CPU_TYPE_I386\n'
+        self.assertEquals(regex.architecture(headers), 'x86_64')
+
+    def test_libswiftcore_x86_64(self):
+        # These are the headers for libswiftCore.dylib when built
+        # for a Darwin x86_64 target.
+        headers = 'Fat headers\n' + \
+                  'fat_magic FAT_MAGIC\n' + \
+                  'nfat_arch 1\n' + \
+                  'architecture x86_64\n' + \
+                  '    cputype CPU_TYPE_X86_64\n' + \
+                  '    cpusubtype CPU_SUBTYPE_X86_64_ALL\n' + \
+                  '    capabilities 0x0\n' + \
+                  '    offset 4096\n' + \
+                  '    size 9029488\n' + \
+                  '    align 2^12 (4096)\n'
+        self.assertEquals(regex.architecture(headers), 'x86_64')