Support constant definition

Thriftpy · Dec 9, 2018 · 91c16b5 · 91c16b5
1 parent 72f3796
commit 91c16b5
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 6 deletions.
diff --git a/tests/recursive_definition.thrift b/tests/recursive_definition.thrift
@@ -4,11 +4,13 @@ service PingPong {
 
 struct Foo {
     1: optional Bar test,
-    2: optional Some some,
+    2: optional SomeInt some_int,
 }
 
 struct Bar {
     1: optional Foo test,
 }
 
-typedef i32 Some
+const SomeInt SOME_INT = [1, 2, 3]
+
+typedef list<i32> SomeInt
diff --git a/tests/test_recursive_definition.py b/tests/test_recursive_definition.py
@@ -7,4 +7,9 @@ def test_recursive_definition():
     thrift = load('./recursive_definition.thrift')
     assert thrift.Bar.thrift_spec == {1: (12, 'test', thrift.Foo, False)}
     assert thrift.Foo.thrift_spec == {
-        1: (12, 'test', thrift.Bar, False), 2: (8, 'some', False)}
+        1: (12, 'test', thrift.Bar, False), 2: (8, 'some_int', False)}
+
+
+def test_const():
+    thrift = load('./recursive_definition.thrift')
+    assert thrift.SOME_INT == [1, 2, 3]
diff --git a/thriftpy/parser/__init__.py b/thriftpy/parser/__init__.py
@@ -13,7 +13,7 @@
 import sys
 import types
 
-from .parser import parse, parse_fp, incomplete_type
+from .parser import parse, parse_fp, incomplete_type, _cast
 from .exc import ThriftParserError
 from ..thrift import TPayloadMeta
 
@@ -41,6 +41,10 @@ def load(path, module_name=None, include_dirs=None, include_dir=None):
 def fill_incomplete_ttype(tmodule, definition):
     # fill incomplete ttype
     if isinstance(definition, tuple):
+        # fill const value
+        if definition[0] == 'UNKNOWN_CONST':
+            ttype = get_definition(tmodule, incomplete_type[definition[1]][0], definition[3])
+            return _cast(ttype)(definition[2])
         # fill an incomplete alias ttype
         if definition[1] in incomplete_type:
             return (definition[0], get_definition(tmodule, *incomplete_type[definition[1]]))

diff --git a/thriftpy/parser/parser.py b/thriftpy/parser/parser.py
@@ -108,7 +108,7 @@ def p_const(p):
              | CONST field_type IDENTIFIER '=' const_value sep'''
 
     try:
-        val = _cast(p[2])(p[5])
+        val = _cast(p[2], p.lineno(3))(p[5])
     except AssertionError:
         raise ThriftParserError('Type error for constant %s at line %d' %
                                 (p[3], p.lineno(3)))
@@ -656,7 +656,9 @@ def _parse_seq(p):
         p[0] = []
 
 
-def _cast(t):  # noqa
+def _cast(t, linno=0):  # noqa
+    if t < 0:
+        return _lazy_cast_const(t, linno)
     if t == TType.BOOL:
         return _cast_bool
     if t == TType.BYTE:
@@ -685,6 +687,12 @@ def _cast(t):  # noqa
         return _cast_struct(t)
 
 
+def _lazy_cast_const(t, linno):
+    def _inner_cast(v):
+        return ('UNKNOWN_CONST', t, v, linno)
+    return _inner_cast
+
+
 def _cast_bool(v):
     assert isinstance(v, (bool, int))
     return bool(v)