Merge pull request #131 from hit9/enable-multiple-include-dirs

Enable multiple include dirs

tests/test_parser.py

@@ -28,12 +28,14 @@ def test_constants():
 
 
 def test_include():
-    thrift = load('parser-cases/include.thrift', include_dir='./parser-cases')
+    thrift = load('parser-cases/include.thrift', include_dirs=[
+        './parser-cases'])
     assert thrift.datetime == 1422009523
 
 
 def test_tutorial():
-    thrift = load('parser-cases/tutorial.thrift', include_dir='./parser-cases')
+    thrift = load('parser-cases/tutorial.thrift', include_dirs=[
+        './parser-cases'])
     assert thrift.INT32CONSTANT == 9853
     assert thrift.MAPCONSTANT == {'hello': 'world', 'goodnight': 'moon'}
     assert thrift.Operation.ADD == 1 and thrift.Operation.SUBTRACT == 2 \

thriftpy/parser/__init__.py

@@ -15,13 +15,18 @@
 from .parser import parse
 
 
-def load(path, module_name=None, include_dir=None):
+def load(path, module_name=None, include_dirs=None, include_dir=None):
     """Load thrift_file as a module
     The module loaded and objects inside may only be pickled if module_name
     was provided.
+
+    Note: `include_dir` will be depreacated in the future, use `include_dirs`
+    instead. If `include_dir` was provided (not None), it will be appended to
+    `include_dirs`.
     """
     real_module = bool(module_name)
-    thrift = parse(path, module_name, include_dir=include_dir)
+    thrift = parse(path, module_name, include_dirs=include_dirs,
+                   include_dir=include_dir)
 
     if real_module:
         sys.modules[module_name] = thrift

thriftpy/parser/parser.py

@@ -44,9 +44,15 @@ def p_header_unit(p):
 def p_include(p):
     '''include : INCLUDE LITERAL'''
     thrift = thrift_stack[-1]
-    path = os.path.join(include_dir_, p[2])
-    child = parse(path)
-    setattr(thrift, child.__name__, child)
+
+    for include_dir in include_dirs_:
+        path = os.path.join(include_dir, p[2])
+        if os.path.exists(path):
+            child = parse(path)
+            setattr(thrift, child.__name__, child)
+            return
+    raise ThriftParserError(('Couldn\'t include thrift %s in any '
+                             'directories provided') % p[2])
 
 
 def p_namespace(p):
@@ -390,11 +396,11 @@ def p_definition_type(p):
 
 
 thrift_stack = []
-include_dir_ = '.'
+include_dirs_ = ['.']
 thrift_cache = {}
 
 
-def parse(path, module_name=None, include_dir=None,
+def parse(path, module_name=None, include_dirs=None, include_dir=None,
           lexer=None, parser=None, enable_cache=True):
 
     # dead include checking on current stack
@@ -414,10 +420,12 @@ def parse(path, module_name=None, include_dir=None,
     if parser is None:
         parser = yacc.yacc(debug=False, write_tables=0)
 
-    global include_dir_
+    global include_dirs_
 
+    if include_dirs is not None:
+        include_dirs_ = include_dirs
     if include_dir is not None:
-        include_dir_ = include_dir
+        include_dirs_.append(include_dir)
 
     if not path.endswith('.thrift'):
         raise ThriftParserError('Path should end with .thrift')