[clang.py] Refactor get_tu and get_cursor test helper functions into util.py

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@152510 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 130 insertions, 112 deletions

diff --git a/bindings/python/tests/cindex/test_cursor.py b/bindings/python/tests/cindex/test_cursor.py
index d58ca95b7e..bbe2c1d7e5 100644
--- a/bindings/python/tests/cindex/test_cursor.py
+++ b/bindings/python/tests/cindex/test_cursor.py
@@ -1,4 +1,7 @@
-from clang.cindex import Index, CursorKind, TypeKind
+from clang.cindex import CursorKind
+from clang.cindex import TypeKind
+from .util import get_cursor
+from .util import get_tu
 
 kInput = """\
 // FIXME: Find nicer way to drop builtins and other cruft.
@@ -24,9 +27,8 @@ void f0(int a0, int a1) {
 """
 
 def test_get_children():
-    index = Index.create()
-    tu = index.parse('t.c', unsaved_files = [('t.c',kInput)])
-    
+    tu = get_tu(kInput)
+
     # Skip until past start_decl.
     it = tu.cursor.get_children()
     while it.next().spelling != 'start_decl':
@@ -65,30 +67,18 @@ def test_get_children():
     assert tu_nodes[2].is_definition() == True
 
 def test_underlying_type():
-    source = 'typedef int foo;'
-    index = Index.create()
-    tu = index.parse('test.c', unsaved_files=[('test.c', source)])
-    assert tu is not None
-
-    for cursor in tu.cursor.get_children():
-        if cursor.spelling == 'foo':
-            typedef = cursor
-            break
+    tu = get_tu('typedef int foo;')
+    typedef = get_cursor(tu, 'foo')
+    assert typedef is not None
 
     assert typedef.kind.is_declaration()
     underlying = typedef.underlying_typedef_type
     assert underlying.kind == TypeKind.INT
 
 def test_enum_type():
-    source = 'enum TEST { FOO=1, BAR=2 };'
-    index = Index.create()
-    tu = index.parse('test.c', unsaved_files=[('test.c', source)])
-    assert tu is not None
-
-    for cursor in tu.cursor.get_children():
-        if cursor.spelling == 'TEST':
-            enum = cursor
-            break
+    tu = get_tu('enum TEST { FOO=1, BAR=2 };')
+    enum = get_cursor(tu, 'TEST')
+    assert enum is not None
 
     assert enum.kind == CursorKind.ENUM_DECL
     enum_type = enum.enum_type
diff --git a/bindings/python/tests/cindex/test_diagnostics.py b/bindings/python/tests/cindex/test_diagnostics.py
index b9872d9ce9..48ab6176fd 100644
--- a/bindings/python/tests/cindex/test_diagnostics.py
+++ b/bindings/python/tests/cindex/test_diagnostics.py
@@ -1,19 +1,10 @@
 from clang.cindex import *
-
-def tu_from_source(source, all_warnings=False):
-    args = []
-    if all_warnings:
-        args = ['-Wall', '-Wextra']
-
-    index = Index.create()
-    tu = index.parse('INPUT.c', args=args,
-            unsaved_files = [('INPUT.c', source)])
-    return tu
+from .util import get_tu
 
 # FIXME: We need support for invalid translation units to test better.
 
 def test_diagnostic_warning():
-    tu = tu_from_source("""int f0() {}\n""")
+    tu = get_tu('int f0() {}\n')
     assert len(tu.diagnostics) == 1
     assert tu.diagnostics[0].severity == Diagnostic.Warning
     assert tu.diagnostics[0].location.line == 1
@@ -23,8 +14,7 @@ def test_diagnostic_warning():
 
 def test_diagnostic_note():
     # FIXME: We aren't getting notes here for some reason.
-    index = Index.create()
-    tu = tu_from_source("""#define A x\nvoid *A = 1;\n""")
+    tu = get_tu('#define A x\nvoid *A = 1;\n')
     assert len(tu.diagnostics) == 1
     assert tu.diagnostics[0].severity == Diagnostic.Warning
     assert tu.diagnostics[0].location.line == 2
@@ -36,8 +26,7 @@ def test_diagnostic_note():
 #    assert tu.diagnostics[1].spelling == 'instantiated from'
 
 def test_diagnostic_fixit():
-    index = Index.create()
-    tu = tu_from_source("""struct { int f0; } x = { f0 : 1 };""")
+    tu = get_tu('struct { int f0; } x = { f0 : 1 };')
     assert len(tu.diagnostics) == 1
     assert tu.diagnostics[0].severity == Diagnostic.Warning
     assert tu.diagnostics[0].location.line == 1
@@ -51,8 +40,7 @@ def test_diagnostic_fixit():
     assert tu.diagnostics[0].fixits[0].value == '.f0 = '
 
 def test_diagnostic_range():
-    index = Index.create()
-    tu = tu_from_source("""void f() { int i = "a" + 1; }""")
+    tu = get_tu('void f() { int i = "a" + 1; }')
     assert len(tu.diagnostics) == 1
     assert tu.diagnostics[0].severity == Diagnostic.Warning
     assert tu.diagnostics[0].location.line == 1
@@ -72,9 +60,8 @@ def test_diagnostic_range():
       assert False
 
 def test_diagnostic_category():
-    # Ensure that category properties work.
-    index = Index.create()
-    tu = tu_from_source("""int f(int i) { return 7; }""", all_warnings=True)
+    """Ensure that category properties work."""
+    tu = get_tu('int f(int i) { return 7; }', all_warnings=True)
     assert len(tu.diagnostics) == 1
     d = tu.diagnostics[0]
 
@@ -86,12 +73,10 @@ def test_diagnostic_category():
     assert d.category_name == 'Semantic Issue'
 
 def test_diagnostic_option():
-    # Ensure that category option properties work.
-    index = Index.create()
-    tu = tu_from_source("""int f(int i) { return 7; }""", all_warnings=True)
+    """Ensure that category option properties work."""
+    tu = get_tu('int f(int i) { return 7; }', all_warnings=True)
     assert len(tu.diagnostics) == 1
     d = tu.diagnostics[0]
 
     assert d.option == '-Wunused-parameter'
     assert d.disable_option == '-Wno-unused-parameter'
-
diff --git a/bindings/python/tests/cindex/test_location.py b/bindings/python/tests/cindex/test_location.py
index 1707f01aea..528676ef14 100644
--- a/bindings/python/tests/cindex/test_location.py
+++ b/bindings/python/tests/cindex/test_location.py
@@ -1,4 +1,9 @@
-from clang.cindex import Index, File, SourceLocation, SourceRange, Cursor
+from clang.cindex import Cursor
+from clang.cindex import File
+from clang.cindex import SourceLocation
+from clang.cindex import SourceRange
+from .util import get_cursor
+from .util import get_tu
 
 baseInput="int one;\nint two;\n"
 
@@ -8,44 +13,46 @@ def assert_location(loc, line, column, offset):
     assert loc.offset == offset
 
 def test_location():
-    index = Index.create()
-    tu = index.parse('t.c', unsaved_files = [('t.c',baseInput)])
+    tu = get_tu(baseInput)
+    one = get_cursor(tu, 'one')
+    two = get_cursor(tu, 'two')
 
-    for n in tu.cursor.get_children():
-        if n.spelling == 'one':
-            assert_location(n.location,line=1,column=5,offset=4)
-        if n.spelling == 'two':
-            assert_location(n.location,line=2,column=5,offset=13)
+    assert one is not None
+    assert two is not None
+
+    assert_location(one.location,line=1,column=5,offset=4)
+    assert_location(two.location,line=2,column=5,offset=13)
 
     # adding a linebreak at top should keep columns same
-    tu = index.parse('t.c', unsaved_files = [('t.c',"\n"+baseInput)])
+    tu = get_tu('\n' + baseInput)
+    one = get_cursor(tu, 'one')
+    two = get_cursor(tu, 'two')
+
+    assert one is not None
+    assert two is not None
 
-    for n in tu.cursor.get_children():
-        if n.spelling == 'one':
-            assert_location(n.location,line=2,column=5,offset=5)
-        if n.spelling == 'two':
-            assert_location(n.location,line=3,column=5,offset=14)
+    assert_location(one.location,line=2,column=5,offset=5)
+    assert_location(two.location,line=3,column=5,offset=14)
 
     # adding a space should affect column on first line only
-    tu = index.parse('t.c', unsaved_files = [('t.c'," "+baseInput)])
+    tu = get_tu(' ' + baseInput)
+    one = get_cursor(tu, 'one')
+    two = get_cursor(tu, 'two')
 
-    for n in tu.cursor.get_children():
-        if n.spelling == 'one':
-            assert_location(n.location,line=1,column=6,offset=5)
-        if n.spelling == 'two':
-            assert_location(n.location,line=2,column=5,offset=14)
+    assert_location(one.location,line=1,column=6,offset=5)
+    assert_location(two.location,line=2,column=5,offset=14)
 
     # define the expected location ourselves and see if it matches
     # the returned location
-    tu = index.parse('t.c', unsaved_files = [('t.c',baseInput)])
+    tu = get_tu(baseInput)
 
     file = File.from_name(tu, 't.c')
     location = SourceLocation.from_position(tu, file, 1, 5)
     cursor = Cursor.from_location(tu, location)
 
-    for n in tu.cursor.get_children():
-        if n.spelling == 'one':
-            assert n == cursor
+    one = get_cursor(tu, 'one')
+    assert one is not None
+    assert one == cursor
 
     # Ensure locations referring to the same entity are equivalent.
     location2 = SourceLocation.from_position(tu, file, 1, 5)
@@ -54,18 +61,17 @@ def test_location():
     assert location2 != location3
 
 def test_extent():
-    index = Index.create()
-    tu = index.parse('t.c', unsaved_files = [('t.c',baseInput)])
-
-    for n in tu.cursor.get_children():
-        if n.spelling == 'one':
-            assert_location(n.extent.start,line=1,column=1,offset=0)
-            assert_location(n.extent.end,line=1,column=8,offset=7)
-            assert baseInput[n.extent.start.offset:n.extent.end.offset] == "int one"
-        if n.spelling == 'two':
-            assert_location(n.extent.start,line=2,column=1,offset=9)
-            assert_location(n.extent.end,line=2,column=8,offset=16)
-            assert baseInput[n.extent.start.offset:n.extent.end.offset] == "int two"
+    tu = get_tu(baseInput)
+    one = get_cursor(tu, 'one')
+    two = get_cursor(tu, 'two')
+
+    assert_location(one.extent.start,line=1,column=1,offset=0)
+    assert_location(one.extent.end,line=1,column=8,offset=7)
+    assert baseInput[one.extent.start.offset:one.extent.end.offset] == "int one"
+
+    assert_location(two.extent.start,line=2,column=1,offset=9)
+    assert_location(two.extent.end,line=2,column=8,offset=16)
+    assert baseInput[two.extent.start.offset:two.extent.end.offset] == "int two"
 
     file = File.from_name(tu, 't.c')
     location1 = SourceLocation.from_position(tu, file, 1, 1)
diff --git a/bindings/python/tests/cindex/test_type.py b/bindings/python/tests/cindex/test_type.py
index e4cd245c9c..ed852103b9 100644
--- a/bindings/python/tests/cindex/test_type.py
+++ b/bindings/python/tests/cindex/test_type.py
@@ -1,8 +1,9 @@
-from clang.cindex import Cursor
 from clang.cindex import CursorKind
 from clang.cindex import Index
 from clang.cindex import TypeKind
 from nose.tools import raises
+from .util import get_cursor
+from .util import get_tu
 
 kInput = """\
 
@@ -21,35 +22,6 @@ struct teststruct {
 
 """
 
-def get_tu(source=kInput, lang='c'):
-    name = 't.c'
-    if lang == 'cpp':
-        name = 't.cpp'
-
-    index = Index.create()
-    tu = index.parse(name, unsaved_files=[(name, source)])
-    assert tu is not None
-    return tu
-
-def get_cursor(source, spelling):
-    children = []
-    if isinstance(source, Cursor):
-        children = source.get_children()
-    else:
-        # Assume TU
-        children = source.cursor.get_children()
-
-    for cursor in children:
-        if cursor.spelling == spelling:
-            return cursor
-
-        # Recurse into children.
-        result = get_cursor(cursor, spelling)
-        if result is not None:
-            return result
-
-    return None
-
 def test_a_struct():
     tu = get_tu(kInput)
 
diff --git a/bindings/python/tests/cindex/util.py b/bindings/python/tests/cindex/util.py
new file mode 100644
index 0000000000..388b269489
--- /dev/null
+++ b/bindings/python/tests/cindex/util.py
@@ -0,0 +1,65 @@
+# This file provides common utility functions for the test suite.
+
+from clang.cindex import Cursor
+from clang.cindex import Index
+
+def get_tu(source, lang='c', all_warnings=False):
+    """Obtain a translation unit from source and language.
+
+    By default, the translation unit is created from source file "t.<ext>"
+    where <ext> is the default file extension for the specified language. By
+    default it is C, so "t.c" is the default file name.
+
+    Supported languages are {c, cpp, objc}.
+
+    all_warnings is a convenience argument to enable all compiler warnings.
+    """
+    name = 't.c'
+    if lang == 'cpp':
+        name = 't.cpp'
+    elif lang == 'objc':
+        name = 't.m'
+    elif lang != 'c':
+        raise Exception('Unknown language: %s' % lang)
+
+    args = []
+    if all_warnings:
+        args = ['-Wall', '-Wextra']
+
+    index = Index.create()
+    tu = index.parse(name, args=args, unsaved_files=[(name, source)])
+    assert tu is not None
+    return tu
+
+def get_cursor(source, spelling):
+    """Obtain a cursor from a source object.
+
+    This provides a convenient search mechanism to find a cursor with specific
+    spelling within a source. The first argument can be either a
+    TranslationUnit or Cursor instance.
+
+    If the cursor is not found, None is returned.
+    """
+    children = []
+    if isinstance(source, Cursor):
+        children = source.get_children()
+    else:
+        # Assume TU
+        children = source.cursor.get_children()
+
+    for cursor in children:
+        if cursor.spelling == spelling:
+            return cursor
+
+        # Recurse into children.
+        result = get_cursor(cursor, spelling)
+        if result is not None:
+            return result
+
+    return None
+
+
+__all__ = [
+    'get_cursor',
+    'get_tu',
+]