[clang.py] Add TranslationUnit.get_{file,source_location,source_range}

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

2 files changed, 127 insertions, 0 deletions

diff --git a/bindings/python/clang/cindex.py b/bindings/python/clang/cindex.py
index 8762417fc4..1bc7d0bae9 100644
--- a/bindings/python/clang/cindex.py
+++ b/bindings/python/clang/cindex.py
@@ -1895,6 +1895,67 @@ class TranslationUnit(ClangObject):
 
         return iter(includes)
 
+    def get_file(self, filename):
+        """Obtain a File from this translation unit."""
+
+        return File.from_name(self, filename)
+
+    def get_location(self, filename, position):
+        """Obtain a SourceLocation for a file in this translation unit.
+
+        The position can be specified by passing:
+
+          - Integer file offset. Initial file offset is 0.
+          - 2-tuple of (line number, column number). Initial file position is
+            (0, 0)
+        """
+        f = self.get_file(filename)
+
+        if isinstance(position, int):
+            return SourceLocation.from_offset(self, f, position)
+
+        return SourceLocation.from_position(self, f, position[0], position[1])
+
+    def get_extent(self, filename, locations):
+        """Obtain a SourceRange from this translation unit.
+
+        The bounds of the SourceRange must ultimately be defined by a start and
+        end SourceLocation. For the locations argument, you can pass:
+
+          - 2 SourceLocation instances in a 2-tuple or list.
+          - 2 int file offsets via a 2-tuple or list.
+          - 2 2-tuple or lists of (line, column) pairs in a 2-tuple or list.
+
+        e.g.
+
+        get_extent('foo.c', (5, 10))
+        get_extent('foo.c', ((1, 1), (1, 15)))
+        """
+        f = self.get_file(filename)
+
+        if len(locations) < 2:
+            raise Exception('Must pass object with at least 2 elements')
+
+        start_location, end_location = locations
+
+        if hasattr(start_location, '__len__'):
+            start_location = SourceLocation.from_position(self, f,
+                start_location[0], start_location[1])
+        elif isinstance(start_location, int):
+            start_location = SourceLocation.from_offset(self, f,
+                start_location)
+
+        if hasattr(end_location, '__len__'):
+            end_location = SourceLocation.from_position(self, f,
+                end_location[0], end_location[1])
+        elif isinstance(end_location, int):
+            end_location = SourceLocation.from_offset(self, f, end_location)
+
+        assert isinstance(start_location, SourceLocation)
+        assert isinstance(end_location, SourceLocation)
+
+        return SourceRange.from_locations(start_location, end_location)
+
     @property
     def diagnostics(self):
         """
diff --git a/bindings/python/tests/cindex/test_translation_unit.py b/bindings/python/tests/cindex/test_translation_unit.py
index 982a608943..9de12ad462 100644
--- a/bindings/python/tests/cindex/test_translation_unit.py
+++ b/bindings/python/tests/cindex/test_translation_unit.py
@@ -1,6 +1,9 @@
 from clang.cindex import CursorKind
 from clang.cindex import Cursor
+from clang.cindex import File
 from clang.cindex import Index
+from clang.cindex import SourceLocation
+from clang.cindex import SourceRange
 from clang.cindex import TranslationUnitSaveError
 from clang.cindex import TranslationUnit
 from .util import get_cursor
@@ -151,3 +154,66 @@ def test_index_parse():
     index = Index.create()
     tu = index.parse(path)
     assert isinstance(tu, TranslationUnit)
+
+def test_get_file():
+    """Ensure tu.get_file() works appropriately."""
+
+    tu = get_tu('int foo();')
+
+    f = tu.get_file('t.c')
+    assert isinstance(f, File)
+    assert f.name == 't.c'
+
+    try:
+        f = tu.get_file('foobar.cpp')
+    except:
+        pass
+    else:
+        assert False
+
+def test_get_source_location():
+    """Ensure tu.get_source_location() works."""
+
+    tu = get_tu('int foo();')
+
+    location = tu.get_location('t.c', 2)
+    assert isinstance(location, SourceLocation)
+    assert location.offset == 2
+    assert location.file.name == 't.c'
+
+    location = tu.get_location('t.c', (1, 3))
+    assert isinstance(location, SourceLocation)
+    assert location.line == 1
+    assert location.column == 3
+    assert location.file.name == 't.c'
+
+def test_get_source_range():
+    """Ensure tu.get_source_range() works."""
+
+    tu = get_tu('int foo();')
+
+    r = tu.get_extent('t.c', (1,4))
+    assert isinstance(r, SourceRange)
+    assert r.start.offset == 1
+    assert r.end.offset == 4
+    assert r.start.file.name == 't.c'
+    assert r.end.file.name == 't.c'
+
+    r = tu.get_extent('t.c', ((1,2), (1,3)))
+    assert isinstance(r, SourceRange)
+    assert r.start.line == 1
+    assert r.start.column == 2
+    assert r.end.line == 1
+    assert r.end.column == 3
+    assert r.start.file.name == 't.c'
+    assert r.end.file.name == 't.c'
+
+    start = tu.get_location('t.c', 0)
+    end = tu.get_location('t.c', 5)
+
+    r = tu.get_extent('t.c', (start, end))
+    assert isinstance(r, SourceRange)
+    assert r.start.offset == 0
+    assert r.end.offset == 5
+    assert r.start.file.name == 't.c'
+    assert r.end.file.name == 't.c'