Enable writing module code directly to file

libcst/_nodes/internal.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import io
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Iterable, Iterator, List, Optional, Sequence, TYPE_CHECKING, Union
@@ -70,6 +71,30 @@
  yield
 
 
+@add_slots
+@dataclass(frozen=False)
+class CodegenWriter(CodegenState):
+ """
+ A CodegenState that writes to a file-like object.
+ """
+
+ writer: io.TextIOBase = None # need a default value for dataclass
+
+ def __post_init__(self) -> None:
+ if self.writer is None:
+ raise TypeError("writer must be provided")
+
+ def add_indent_tokens(self) -> None:
+ for token in self.indent_tokens:
+ self.writer.write(token)
+
+ def add_token(self, value: str) -> None:
+ self.writer.write(value)
+
+ def pop_trailing_newline(self) -> None:
+ pass
+
+
 def visit_required(
  parent: "CSTNode", fieldname: str, node: CSTNodeT, visitor: "CSTVisitorT"
 ) -> CSTNodeT:

libcst/_nodes/module.py

@@ -3,12 +3,18 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import io
 from dataclasses import dataclass
 from typing import cast, Optional, Sequence, TYPE_CHECKING, TypeVar, Union
 
 from libcst._add_slots import add_slots
 from libcst._nodes.base import CSTNode
-from libcst._nodes.internal import CodegenState, visit_body_sequence, visit_sequence
+from libcst._nodes.internal import (
+ CodegenState,
+ CodegenWriter,
+ visit_body_sequence,
+ visit_sequence,
+)
 from libcst._nodes.statement import (
  BaseCompoundStatement,
  get_docstring_impl,
@@ -136,6 +142,24 @@
  node._codegen(state)
  return "".join(state.tokens)
 
+ def write_code(self, writer: io.TextIOBase):
+ """
+ Like :meth:`code`, but writes the code to the given file-like object.
+ """
+ self.write_code_for_node(self, writer)
+
+ def write_code_for_node(self, node: CSTNode, writer: io.TextIOBase):
+ """
+ Like :meth:`code_for_node`, but writes the code to the given file-like object.
+ """
+
+ state = CodegenWriter(
+ default_indent=self.default_indent,
+ default_newline=self.default_newline,
+ writer=writer,
+ )
+ node._codegen(state)
+
  @property
  def config_for_parsing(self) -> "PartialParserConfig":
  """