Repository: dropbox/stone

Issue:
`python_type_stubs` generator incorrectly puts imports at bottom of file, not top
Hi all! This is actually a bug I introduced, just making sure it's tracked somewhere.

# The issue
Mypy *will not* correctly detect the type of `a` in the following code snippet of a `.pyi` file.
```
a = ...  # type: Text
from typing import Text
```
Mypy *will* correctly detect the type of `a` if we literally just swap those lines.
```
from typing import Text
a = ...  # type: Text
```

Unfortunately, my `python_type_stubs` generator does exactly the incorrect behavior, as shown in this test suite:
https://github.com/dropbox/stone/blob/4f4e5906a2be1431ea6f12726959f456d0670086/test/test_python_type_stubs.py#L306

The result: Our apiv2 stuff is often poorly typed as `typing.Any`s :(

# Where does the issue occur?
If we look at  [_generate_base_namespace_module](https://github.com/dropbox/stone/blob/4f4e5906a2be1431ea6f12726959f456d0670086/stone/backends/python_type_stubs.py#L113), sure enough the `_generate_imports_needed_for_typing` is the final thing called. This is because we build up a list of all the things that, well, _do_ need to be imported.

# The solution
Just my opinion, feel free to come up with another solution:

`self.emit_raw` should not immediately write to a file - perhaps split this into something like `self.append_to_output_buffer(...)`, that doesn't actually write anything to a file yet, and a `self.emit_output_buffer(...)`. This would let us do something like:

```
<stone doing its' thing, doing self.append_output_buffer() while figuring out what imports are needed>
self.prepend_output_buffer(the imports required)
self.emit_output_buffer
```

Discussion:

Patch:
diff --git a/stone/backend.py b/stone/backend.py
index 1831546..9f61f37 100644
--- a/stone/backend.py
+++ b/stone/backend.py
@@ -123,6 +123,8 @@ class Backend(object):
         self.output = []  # type: typing.List[typing.Text]
         self.lineno = 1
         self.cur_indent = 0
+        self.positional_placeholders = []  # type: typing.List[typing.Text]
+        self.named_placeholders = {}  # type: typing.Dict[typing.Text, typing.Text]
 
         self.args = None  # type: typing.Optional[argparse.Namespace]
 
@@ -165,19 +167,23 @@ class Backend(object):
             os.makedirs(directory)
 
         self.logger.info('Generating %s', full_path)
-        self.output = []
+        self.clear_output_buffer()
         yield
         with open(full_path, 'wb') as f:
-            f.write(''.join(self.output).encode('utf-8'))
-        self.output = []
+            f.write(self.output_buffer_to_string().encode('utf-8'))
+        self.clear_output_buffer()
 
     def output_buffer_to_string(self):
         # type: () -> typing.Text
         """Returns the contents of the output buffer as a string."""
-        return ''.join(self.output)
+        return ''.join(self.output).format(
+            *self.positional_placeholders,
+            **self.named_placeholders)
 
     def clear_output_buffer(self):
         self.output = []
+        self.positional_placeholders = []
+        self.named_placeholders = {}
 
     def indent_step(self):
         # type: () -> int
@@ -214,6 +220,15 @@ class Backend(object):
         else:
             return ' ' * self.cur_indent
 
+    @contextmanager
+    def capture_emitted_output(self, output_buffer):
+        # type: (six.StringIO) -> typing.Iterator[None]
+        original_output = self.output
+        self.output = []
+        yield
+        output_buffer.write(''.join(self.output))
+        self.output = original_output
+
     def emit_raw(self, s):
         # type: (typing.Text) -> None
         """
@@ -222,7 +237,7 @@ class Backend(object):
         indentation is generated.
         """
         self.lineno += s.count('\n')
-        self._append_output(s)
+        self._append_output(s.replace('{', '{{').replace('}', '}}'))
         if len(s) > 0 and s[-1] != '\n':
             raise AssertionError(
                 'Input string to emit_raw must end with a newline.')
@@ -291,6 +306,27 @@ class Backend(object):
                                     break_on_hyphens=break_on_hyphens,
                                     ) + '\n')
 
+    def emit_placeholder(self, s=''):
+        # type: (typing.Text) -> None
+        """
+        Emits replacements fields that can be used to format the output string later.
+        """
+        self._append_output('{%s}' % s)
+
+    def add_positional_placeholder(self, s):
+        # type: (typing.Text) -> None
+        """
+        Format replacement fields corresponding to empty calls to emit_placeholder.
+        """
+        self.positional_placeholders.append(s)
+
+    def add_named_placeholder(self, name, s):
+        # type: (typing.Text, typing.Text) -> None
+        """
+        Format replacement fields corresponding to non-empty calls to emit_placeholder.
+        """
+        self.named_placeholders[name] = s
+
     @classmethod
     def process_doc(cls, doc, handler):
         # type: (str, typing.Callable[[str, str], str]) -> typing.Text
diff --git a/stone/backends/js_client.py b/stone/backends/js_client.py
index d281071..96a87aa 100644
--- a/stone/backends/js_client.py
+++ b/stone/backends/js_client.py
@@ -102,7 +102,7 @@ class JavascriptClientBackend(CodeBackend):
                     'return this.request("%s", arg);' % url)
         self.emit('};')
 
-    def _docf(self, tag, val):  # pylint: disable=unused-argument
+    def _docf(self, tag, val):
         """
         Callback used as the handler argument to process_docs(). This converts
         Stone doc references to JSDoc-friendly annotations.
diff --git a/stone/backends/python_type_stubs.py b/stone/backends/python_type_stubs.py
index 6f38967..ce532c2 100644
--- a/stone/backends/python_type_stubs.py
+++ b/stone/backends/python_type_stubs.py
@@ -1,5 +1,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+from six import StringIO
+
 import textwrap
 
 from stone.backend import CodeBackend
@@ -26,6 +28,7 @@ from stone.ir import (  # noqa: F401 # pylint: disable=unused-import
     ApiNamespace,
     DataType,
     is_nullable_type,
+    is_struct_type,
     is_union_type,
     is_user_defined_type,
     is_void_type,
@@ -112,6 +115,7 @@ class PythonTypeStubsBackend(CodeBackend):
         self.import_tracker.clear()
         generate_module_header(self)
 
+        self.emit_placeholder('imports_needed_for_typing')
         self.emit_raw(validators_import_with_type_ignore)
 
         # Generate import statements for all referenced namespaces.
@@ -166,21 +170,21 @@ class PythonTypeStubsBackend(CodeBackend):
 
     def _generate_annotation_type_class_init(self, ns, annotation_type):
         # type: (ApiNamespace, AnnotationType) -> None
-        args = []
+        args = ['self']
         for param in annotation_type.params:
             param_name = fmt_var(param.name, True)
             param_type = self.map_stone_type_to_pep484_type(ns, param.data_type)
+
+            if not is_nullable_type(param.data_type):
+                self.import_tracker._register_typing_import('Optional')
+                param_type = 'Optional[{}]'.format(param_type)
+
             args.append(
                 "{param_name}: {param_type} = ...".format(
                     param_name=param_name,
-                    param_type=param_type,
-                )
-            )
+                    param_type=param_type))
 
-        self.generate_multiline_list(
-            before='def __init__',
-            items=["self"] + args,
-            after=' -> None: ...')
+        self.generate_multiline_list(args, before='def __init__', after=' -> None: ...')
         self.emit()
 
     def _generate_annotation_type_class_properties(self, ns, annotation_type):
@@ -313,8 +317,10 @@ class PythonTypeStubsBackend(CodeBackend):
         if data_type.parent_type:
             extends = class_name_for_data_type(data_type.parent_type, ns)
         else:
-            if is_union_type(data_type):
+            if is_struct_type(data_type):
                 # Use a handwritten base class
+                extends = 'bb.Struct'
+            elif is_union_type(data_type):
                 extends = 'bb.Union'
             else:
                 extends = 'object'
@@ -323,35 +329,28 @@ class PythonTypeStubsBackend(CodeBackend):
 
     def _generate_struct_class_init(self, ns, struct):
         # type: (ApiNamespace, Struct) -> None
-        args = []
+        args = ["self"]
         for field in struct.all_fields:
             field_name_reserved_check = fmt_var(field.name, True)
             field_type = self.map_stone_type_to_pep484_type(ns, field.data_type)
 
-            # The constructor takes all possible fields as optional arguments
-            if not is_nullable_type(field.data_type) and not is_void_type(field.data_type):
+            if not is_nullable_type(field.data_type):
                 self.import_tracker._register_typing_import('Optional')
                 field_type = 'Optional[{}]'.format(field_type)
 
-            args.append(
-                "{field_name}: {field_type} = ...".format(
-                    field_name=field_name_reserved_check,
-                    field_type=field_type,
-                )
-            )
+            args.append("{field_name}: {field_type} = ...".format(
+                field_name=field_name_reserved_check,
+                field_type=field_type))
 
-        self.generate_multiline_list(
-            before='def __init__',
-            items=["self"] + args,
-            after=' -> None: ...')
+        self.generate_multiline_list(args, before='def __init__', after=' -> None: ...')
 
     property_template = textwrap.dedent(
         """
         @property
-        def {field_name}(self) -> {getter_field_type}: ...
+        def {field_name}(self) -> {field_type}: ...
 
         @{field_name}.setter
-        def {field_name}(self, val: {setter_field_type}) -> None: ...
+        def {field_name}(self, val: {field_type}) -> None: ...
 
         @{field_name}.deleter
         def {field_name}(self) -> None: ...
@@ -362,14 +361,12 @@ class PythonTypeStubsBackend(CodeBackend):
         to_emit = []  # type: typing.List[typing.Text]
         for field in struct.all_fields:
             field_name_reserved_check = fmt_func(field.name, check_reserved=True)
-            setter_field_type = self.map_stone_type_to_pep484_type(ns, field.data_type)
-            getter_field_type = setter_field_type
+            field_type = self.map_stone_type_to_pep484_type(ns, field.data_type)
 
             to_emit.extend(
                 self.property_template.format(
                     field_name=field_name_reserved_check,
-                    getter_field_type=getter_field_type,
-                    setter_field_type=setter_field_type
+                    field_type=field_type,
                 ).split("\n")
             )
 
@@ -471,15 +468,19 @@ class PythonTypeStubsBackend(CodeBackend):
 
     def _generate_imports_needed_for_typing(self):
         # type: () -> None
-        if self.import_tracker.cur_namespace_typing_imports:
-            self.emit("")
-            self.emit('from typing import (')
-            with self.indent():
-                for to_import in sorted(self.import_tracker.cur_namespace_typing_imports):
-                    self.emit("{},".format(to_import))
-            self.emit(')')
-
-        if self.import_tracker.cur_namespace_adhoc_imports:
-            self.emit("")
-            for to_import in self.import_tracker.cur_namespace_adhoc_imports:
-                self.emit(to_import)
+        output_buffer = StringIO()
+        with self.capture_emitted_output(output_buffer):
+            if self.import_tracker.cur_namespace_typing_imports:
+                self.emit("")
+                self.emit('from typing import (')
+                with self.indent():
+                    for to_import in sorted(self.import_tracker.cur_namespace_typing_imports):
+                        self.emit("{},".format(to_import))
+                self.emit(')')
+
+            if self.import_tracker.cur_namespace_adhoc_imports:
+                self.emit("")
+                for to_import in self.import_tracker.cur_namespace_adhoc_imports:
+                    self.emit(to_import)
+
+        self.add_named_placeholder('imports_needed_for_typing', output_buffer.getvalue())
diff --git a/stone/backends/python_types.py b/stone/backends/python_types.py
index 4c2719d..54d6dd1 100644
--- a/stone/backends/python_types.py
+++ b/stone/backends/python_types.py
@@ -530,6 +530,7 @@ class PythonTypesBackend(CodeBackend):
         for field in data_type.all_fields:
             field_name_reserved_check = fmt_var(field.name, True)
             args.append('%s=None' % field_name_reserved_check)
+
         self.generate_multiline_list(args, before='def __init__', after=':')
 
         with self.indent():


Test patch:
diff --git a/test/backend_test_util.py b/test/backend_test_util.py
index e824cdc..0b28a44 100644
--- a/test/backend_test_util.py
+++ b/test/backend_test_util.py
@@ -11,18 +11,25 @@ except ImportError:
 
 from stone.backend import Backend  # noqa: F401 # pylint: disable=unused-import
 
-def _mock_emit(backend):
-    # type: (Backend) -> typing.List[str]
+def _mock_output(backend):
+    # type: (Backend) -> typing.Callable[[], str]
     """
     Mock out Backend's .emit function, and return a list containing all params
     emit was called with.
     """
-    recorded_emits = []  # type: typing.List[str]
+    recorded_output = []  # type: typing.List[str]
 
-    def record_emit(s):
-        recorded_emits.append(s)
+    output_buffer_to_string = backend.output_buffer_to_string
 
-    orig_append = backend._append_output
-    backend._append_output = Mock(wraps=orig_append, side_effect=record_emit)  # type: ignore
+    def record_output():
+        recorded_output.append(output_buffer_to_string())
 
-    return recorded_emits
+    backend.output_buffer_to_string = Mock(  # type: ignore
+        wraps=output_buffer_to_string,
+        side_effect=record_output)
+
+    def get_output():
+        backend.output_buffer_to_string = output_buffer_to_string
+        return recorded_output[0] if recorded_output else ''
+
+    return get_output
diff --git a/test/test_js_client.py b/test/test_js_client.py
index 283c74b..9fa9c57 100644
--- a/test/test_js_client.py
+++ b/test/test_js_client.py
@@ -2,7 +2,7 @@ import textwrap
 import unittest
 
 from stone.backends.js_client import JavascriptClientBackend
-from test.backend_test_util import _mock_emit
+from test.backend_test_util import _mock_output
 from stone.ir import Api, ApiNamespace, ApiRoute, Void, Int32
 from stone.ir.data_types import Struct
 
@@ -35,9 +35,9 @@ class TestGeneratedJSClient(unittest.TestCase):
         backend = JavascriptClientBackend(
             target_folder_path='output',
             args=['files', '-c', 'DropboxBase'])
-        emitted = _mock_emit(backend)
+        get_result = _mock_output(backend)
         backend.generate(api)
-        result = "".join(emitted)
+        result = get_result()
 
         expected = textwrap.dedent('''\
             // Auto-generated by Stone, do not modify.
diff --git a/test/test_python_client.py b/test/test_python_client.py
index 41a020b..c6c53dc 100644
--- a/test/test_python_client.py
+++ b/test/test_python_client.py
@@ -2,7 +2,6 @@ import textwrap
 
 from stone.backends.python_client import PythonClientBackend
 from stone.ir import ApiNamespace, ApiRoute, Void, Int32
-from test.backend_test_util import _mock_emit
 
 MYPY = False
 if MYPY:
@@ -20,10 +19,8 @@ class TestGeneratedPythonClient(unittest.TestCase):
         backend = PythonClientBackend(
             target_folder_path='output',
             args=['-m', 'files', '-c', 'DropboxBase', '-t', 'dropbox'])
-        emitted = _mock_emit(backend)
         backend._generate_routes(ns)
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def _evaluate_namespace_with_auth_mode(self, ns, auth_mode):
         # type: (ApiNamespace, str) -> typing.Text
@@ -32,10 +29,8 @@ class TestGeneratedPythonClient(unittest.TestCase):
         backend = PythonClientBackend(
             target_folder_path='output',
             args=['-w', auth_mode, '-m', 'files', '-c', 'DropboxBase', '-t', 'dropbox'])
-        emitted = _mock_emit(backend)
         backend._generate_route_methods({ns})
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def test_route_with_version_number(self):
         # type: () -> None
diff --git a/test/test_python_type_stubs.py b/test/test_python_type_stubs.py
index 57525e6..8ea333d 100644
--- a/test/test_python_type_stubs.py
+++ b/test/test_python_type_stubs.py
@@ -33,7 +33,7 @@ from stone.ir import (
     Float64)
 from stone.ir.api import ApiRoute
 from stone.backends.python_type_stubs import PythonTypeStubsBackend
-from test.backend_test_util import _mock_emit
+
 
 ISO_8601_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
 
@@ -220,6 +220,8 @@ _headers = """\
 # @generated
 # flake8: noqa
 # pylint: skip-file
+
+{}
 try:
     from . import stone_validators as bv
     from . import stone_base as bb
@@ -240,11 +242,8 @@ class TestPythonTypeStubs(unittest.TestCase):
     def _evaluate_namespace(self, ns):
         # type: (ApiNamespace) -> typing.Text
         backend = _make_backend()
-        emitted = _mock_emit(backend)
         backend._generate_base_namespace_module(ns)
-
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def test__generate_base_namespace_module__with_many_structs(self):
         # type: () -> None
@@ -253,7 +252,7 @@ class TestPythonTypeStubs(unittest.TestCase):
         expected = textwrap.dedent("""\
             {headers}
 
-            class Struct1(object):
+            class Struct1(bb.Struct):
                 def __init__(self,
                              f1: Optional[bool] = ...) -> None: ...
 
@@ -274,7 +273,7 @@ class TestPythonTypeStubs(unittest.TestCase):
 
             Struct1_validator: bv.Validator = ...
 
-            class Struct2(object):
+            class Struct2(bb.Struct):
                 def __init__(self,
                              f2: Optional[List[int]] = ...,
                              f3: Optional[datetime.datetime] = ...,
@@ -317,19 +316,18 @@ class TestPythonTypeStubs(unittest.TestCase):
 
             Struct2_validator: bv.Validator = ...
 
-
-            from typing import (
-                Callable,
-                Dict,
-                List,
-                Optional,
-                Text,
-                Type,
-                TypeVar,
-            )
-
-            import datetime
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    Callable,
+                    Dict,
+                    List,
+                    Optional,
+                    Text,
+                    Type,
+                    TypeVar,
+                )
+
+                import datetime""")))
         self.assertEqual(result, expected)
 
     def test__generate_base_namespace_module__with_nested_types(self):
@@ -339,7 +337,7 @@ class TestPythonTypeStubs(unittest.TestCase):
         expected = textwrap.dedent("""\
             {headers}
 
-            class NestedTypes(object):
+            class NestedTypes(bb.Struct):
                 def __init__(self,
                              list_of_nullables: Optional[List[Optional[int]]] = ...,
                              nullable_list: Optional[List[int]] = ...) -> None: ...
@@ -371,15 +369,14 @@ class TestPythonTypeStubs(unittest.TestCase):
 
             NestedTypes_validator: bv.Validator = ...
 
-
-            from typing import (
-                Callable,
-                List,
-                Optional,
-                Type,
-                TypeVar,
-            )
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    Callable,
+                    List,
+                    Optional,
+                    Type,
+                    TypeVar,
+                )""")))
         self.assertEqual(result, expected)
 
     def test__generate_base_namespace_module_with_union__generates_stuff(self):
@@ -425,13 +422,12 @@ class TestPythonTypeStubs(unittest.TestCase):
 
             Shape_validator: bv.Validator = ...
 
-
-            from typing import (
-                Callable,
-                Type,
-                TypeVar,
-            )
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    Callable,
+                    Type,
+                    TypeVar,
+                )""")))
         self.assertEqual(result, expected)
 
     def test__generate_base_namespace_module_with_empty_union(self):
@@ -450,13 +446,12 @@ class TestPythonTypeStubs(unittest.TestCase):
 
             EmptyUnion_validator: bv.Validator = ...
 
-
-            from typing import (
-                Callable,
-                Type,
-                TypeVar,
-            )
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    Callable,
+                    Type,
+                    TypeVar,
+                )""")))
         self.assertEqual(result, expected)
 
     def test__generate_routes(self):
@@ -469,11 +464,10 @@ class TestPythonTypeStubs(unittest.TestCase):
             route_one: bb.Route = ...
             route_one_v2: bb.Route = ...
 
-
-            from typing import (
-                TypeVar,
-            )
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    TypeVar,
+                )""")))
         self.assertEqual(result, expected)
 
     def test__generate_routes_name_conflict(self):
@@ -493,7 +487,7 @@ class TestPythonTypeStubs(unittest.TestCase):
         expected = textwrap.dedent("""\
             {headers}
 
-            class Struct1(object):
+            class Struct1(bb.Struct):
                 def __init__(self,
                              f1: Optional[bool] = ...) -> None: ...
 
@@ -517,12 +511,11 @@ class TestPythonTypeStubs(unittest.TestCase):
             AliasToStruct1_validator: bv.Validator = ...
             AliasToStruct1 = Struct1
             NotUserDefinedAlias_validator: bv.Validator = ...
-
-            from typing import (
-                Callable,
-                Optional,
-                Type,
-                TypeVar,
-            )
-            """).format(headers=_headers)
+            """).format(headers=_headers.format(textwrap.dedent("""\
+                from typing import (
+                    Callable,
+                    Optional,
+                    Type,
+                    TypeVar,
+                )""")))
         self.assertEqual(result, expected)
diff --git a/test/test_python_types.py b/test/test_python_types.py
index 5c00cc2..34ef57b 100644
--- a/test/test_python_types.py
+++ b/test/test_python_types.py
@@ -12,7 +12,6 @@ from stone.ir import (
     StructField,
     Void,
 )
-from test.backend_test_util import _mock_emit
 
 MYPY = False
 if MYPY:
@@ -30,34 +29,29 @@ class TestGeneratedPythonTypes(unittest.TestCase):
         return s
 
     def _mock_backend(self):
-        # type: () -> typing.Tuple[PythonTypesBackend, typing.List]
-        backend = PythonTypesBackend(
+        # type: () -> PythonTypesBackend
+        return PythonTypesBackend(
             target_folder_path='output',
             args=['-r', 'dropbox.dropbox.Dropbox.{ns}_{route}'])
-        emitted = _mock_emit(backend)
-        return backend, emitted
 
     def _evaluate_namespace(self, ns):
         # type: (ApiNamespace) -> typing.Text
-        backend, emitted = self._mock_backend()
+        backend = self._mock_backend()
         route_schema = self._mk_route_schema()
         backend._generate_routes(route_schema, ns)
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def _evaluate_struct(self, ns, struct):
         # type: (ApiNamespace, Struct) -> typing.Text
-        backend, emitted = self._mock_backend()
+        backend = self._mock_backend()
         backend._generate_struct_class(ns, struct)
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def _evaluate_annotation_type(self, ns, annotation_type):
         # type: (ApiNamespace, AnnotationType) -> typing.Text
-        backend, emitted = self._mock_backend()
+        backend = self._mock_backend()
         backend._generate_annotation_type_class(ns, annotation_type)
-        result = "".join(emitted)
-        return result
+        return backend.output_buffer_to_string()
 
     def test_route_with_version_number(self):
         # type: () -> None
@@ -219,7 +213,7 @@ class TestGeneratedPythonTypes(unittest.TestCase):
             MyStruct_validator = bv.Struct(MyStruct)
 
         ''') # noqa
-
+        self.maxDiff = None
         self.assertEqual(result, expected)
 
     def test_annotation_type_class(self):
diff --git a/test/test_tsd_client.py b/test/test_tsd_client.py
index e32e44f..19d105b 100644
--- a/test/test_tsd_client.py
+++ b/test/test_tsd_client.py
@@ -3,7 +3,6 @@ import unittest
 
 from stone.backends.tsd_client import TSDClientBackend
 from stone.ir import Api, ApiNamespace, ApiRoute, Void, Int32
-from test.backend_test_util import _mock_emit
 from stone.ir.data_types import Struct
 
 MYPY = False
@@ -36,9 +35,8 @@ class TestGeneratedTSDClient(unittest.TestCase):
             target_folder_path="output",
             args=['files', 'files']
         )
-        emitted = _mock_emit(backend)
         backend._generate_routes(api, 0, 0)
-        result = "".join(emitted)
+        result = backend.output_buffer_to_string()
         expected = textwrap.dedent(
             '''\
 
diff --git a/test/test_tsd_types.py b/test/test_tsd_types.py
index 9bb34c1..d1cb453 100644
--- a/test/test_tsd_types.py
+++ b/test/test_tsd_types.py
@@ -24,7 +24,7 @@ from stone.ir import (
     Struct,
     StructField)
 from stone.backends.tsd_types import TSDTypesBackend
-from test.backend_test_util import _mock_emit
+from test.backend_test_util import _mock_output
 
 
 def _make_backend(target_folder_path, template_path):
@@ -57,7 +57,7 @@ def _make_struct(struct_name, struct_field_name, namespace):
 def _evaluate_namespace(backend, namespace_list):
     # type: (TSDTypesBackend, typing.List[ApiNamespace]) -> typing.Text
 
-    emitted = _mock_emit(backend)
+    get_result = _mock_output(backend)
     filename = "types.d.ts"
     backend.split_by_namespace = False
     backend._generate_base_namespace_module(namespace_list=namespace_list,
@@ -65,8 +65,7 @@ def _evaluate_namespace(backend, namespace_list):
                                             extra_args={},
                                             template="""/*TYPES*/""",
                                             exclude_error_types=True)
-    result = "".join(emitted)
-    return result
+    return get_result()
 
 
 class TestTSDTypes(unittest.TestCase):