Merge pull request #154 from hit9/load_fp

Add api load_fp to load thrift from file like object

tests/test_parser.py

@@ -2,7 +2,7 @@
 
 import pytest
 from thriftpy.thrift import TType
-from thriftpy.parser import load
+from thriftpy.parser import load, load_fp
 from thriftpy.parser.exc import ThriftParserError, ThriftGrammerError
 
 
@@ -221,3 +221,21 @@ def test_thrift_meta():
     assert meta['exceptions'] == [thrift.InvalidOperation]
     assert meta['services'] == [thrift.Calculator]
     assert meta['includes'] == [thrift.shared]
+
+
+def test_load_fp():
+    thrift = None
+    with open('parser-cases/shared.thrift') as thrift_fp:
+        thrift = load_fp(thrift_fp, 'shared_thrift')
+    assert thrift.__name__ == 'shared_thrift'
+    assert thrift.__thrift_file__ is None
+    assert thrift.__thrift_meta__['structs'] == [thrift.SharedStruct]
+    assert thrift.__thrift_meta__['services'] == [thrift.SharedService]
+
+
+def test_e_load_fp():
+    with pytest.raises(ThriftParserError) as excinfo:
+        with open('parser-cases/tutorial.thrift') as thrift_fp:
+            load_fp(thrift_fp, 'tutorial_thrift')
+        assert ('Unexcepted include statement while loading'
+                'from file like object.') == str(excinfo.value)

thriftpy/__init__.py

@@ -3,8 +3,9 @@
 import sys
 
 from .hook import install_import_hook, remove_import_hook
-from .parser import load, load_module
+from .parser import load, load_module, load_fp
 
 __version__ = '0.3.1'
 __python__ = sys.version_info
-__all__ = ["install_import_hook", "remove_import_hook",  "load", "load_module"]
+__all__ = ["install_import_hook", "remove_import_hook",  "load", "load_module",
+           "load_fp"]

thriftpy/parser/__init__.py

@@ -12,11 +12,12 @@
 import os
 import sys
 
-from .parser import parse
+from .parser import parse, parse_fp
 
 
 def load(path, module_name=None, include_dirs=None, include_dir=None):
-    """Load thrift_file as a module
+    """Load thrift file as a module.
+
     The module loaded and objects inside may only be pickled if module_name
     was provided.
 
@@ -33,6 +34,14 @@ def load(path, module_name=None, include_dirs=None, include_dir=None):
     return thrift
 
 
+def load_fp(source, module_name):
+    """Load thrift file like object as a module.
+    """
+    thrift = parse_fp(source, module_name)
+    sys.modules[module_name] = thrift
+    return thrift
+
+
 def _import_module(import_name):
     if '.' in import_name:
         module, obj = import_name.rsplit('.', 1)

thriftpy/parser/parser.py

@@ -46,6 +46,10 @@ def p_include(p):
     '''include : INCLUDE LITERAL'''
     thrift = thrift_stack[-1]
 
+    if thrift.__thrift_file__ is None:
+        raise ThriftParserError('Unexcepted include statement while loading'
+                                'from file like object.')
+
     for include_dir in include_dirs_:
         path = os.path.join(include_dir, p[2])
         if os.path.exists(path):
@@ -415,10 +419,32 @@ def p_definition_type(p):
 
 def parse(path, module_name=None, include_dirs=None, include_dir=None,
           lexer=None, parser=None, enable_cache=True):
+    """Parse a single thrift file to module object, e.g.::
+
+        >>> from thriftpy.parser.parser import parse
+        >>> note_thrift = parse("path/to/note.thrift")
+        <module 'note_thrift' (built-in)>
+
+    :param path: file path to parse, should be a string ending with '.thrift'.
+    :param module_name: the name for parsed module, the default is the basename
+                        without extension of `path`.
+    :param include_dirs: directories to find thrift files while processing
+                         the `include` directive, by default: ['.'].
+    :param include_dir: directory to find child thrift files. Note this keyword
+                        parameter will be deprecated in the future, it exists
+                        for compatiable reason. If it's provided (not `None`), 
+                        it will be appended to `include_dirs`.
+    :param lexer: ply lexer to use, if not provided, `parse` will new one.
+    :param parser: ply parser to use, if not provided, `parse` will new one.
+    :param enable_cache: if this is set to be `True`, parsed module will be
+                         cached, this is enabled by default. If `module_name`
+                         is provided, use it as cache key, else use the `path`.
+    """
 
     # dead include checking on current stack
     for thrift in thrift_stack:
-        if os.path.samefile(path, thrift.__thrift_file__):
+        if thrift.__thrift_file__ is not None and \
+                os.path.samefile(path, thrift.__thrift_file__):
             raise ThriftParserError('Dead including on %s' % path)
 
     global thrift_cache
@@ -466,6 +492,52 @@ def parse(path, module_name=None, include_dirs=None, include_dir=None,
     return thrift
 
 
+def parse_fp(source, module_name, lexer=None, parser=None, enable_cache=True):
+    """Parse a file-like object to thrift module object, e.g.::
+
+        >>> from thriftpy.parser.parser import parse_fp
+        >>> with open("path/to/note.thrift") as fp:
+                parse_fp(fp, "note_thrift")
+        <module 'note_thrift' (built-in)>
+
+    :param source: file-like object, expected to have a method named `read`.
+    :param module_name: the name for parsed module, shoule be endswith
+                        '_thrift'.
+    :param lexer: ply lexer to use, if not provided, `parse` will new one.
+    :param parser: ply parser to use, if not provided, `parse` will new one.
+    :param enable_cache: if this is set to be `True`, parsed module will be
+                         cached by `module_name`, this is enabled by default.
+    """
+    if not module_name.endswith('_thrift'):
+        raise ThriftParserError('ThriftPy can only generate module with '
+                                '\'_thrift\' suffix')
+
+    if enable_cache and module_name in thrift_cache:
+        return thrift_cache[module_name]
+
+    if not hasattr(source, 'read'):
+        raise ThriftParserError('Except `source` to be a file-like object with'
+                                'a method named \'read\'')
+
+    if lexer is None:
+        lexer = lex.lex()
+    if parser is None:
+        parser = yacc.yacc(debug=False, write_tables=0)
+
+    data = source.read()
+
+    thrift = types.ModuleType(module_name)
+    setattr(thrift, '__thrift_file__', None)
+    thrift_stack.append(thrift)
+    lexer.lineno = 1
+    parser.parse(data)
+    thrift_stack.pop()
+
+    if enable_cache:
+        thrift_cache[module_name] = thrift
+    return thrift
+
+
 def _add_thrift_meta(key, val):
     thrift = thrift_stack[-1]