libcst_native/nodes/
module.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2//
3// This source code is licensed under the MIT license found in the
4// LICENSE file in the root directory of this source tree.
5
6use std::mem::swap;
7
8use crate::tokenizer::whitespace_parser::parse_empty_lines;
9use crate::tokenizer::Token;
10use crate::{
11    nodes::{
12        codegen::{Codegen, CodegenState},
13        statement::*,
14        whitespace::EmptyLine,
15    },
16    tokenizer::whitespace_parser::Config,
17};
18use libcst_derive::cst_node;
19#[cfg(feature = "py")]
20use libcst_derive::TryIntoPy;
21
22use super::traits::{Inflate, Result, WithLeadingLines};
23
24type TokenRef<'r, 'a> = &'r Token<'a>;
25
26#[cst_node]
27pub struct Module<'a> {
28    pub body: Vec<Statement<'a>>,
29    pub header: Vec<EmptyLine<'a>>,
30    pub footer: Vec<EmptyLine<'a>>,
31
32    pub default_indent: &'a str,
33    pub default_newline: &'a str,
34    pub has_trailing_newline: bool,
35    pub encoding: String,
36
37    pub(crate) eof_tok: TokenRef<'a>,
38}
39
40impl<'a> Codegen<'a> for Module<'a> {
41    fn codegen(&self, state: &mut CodegenState<'a>) {
42        for h in &self.header {
43            h.codegen(state);
44        }
45        for s in &self.body {
46            s.codegen(state);
47        }
48        for nl in &self.footer {
49            nl.codegen(state);
50        }
51    }
52}
53
54impl<'r, 'a> Inflate<'a> for DeflatedModule<'r, 'a> {
55    type Inflated = Module<'a>;
56    fn inflate(self, config: &Config<'a>) -> Result<Self::Inflated> {
57        let default_indent = config.default_indent;
58        let default_newline = config.default_newline;
59        let has_trailing_newline = config.has_trailing_newline();
60        let mut body = self.body.inflate(config)?;
61        let mut footer = parse_empty_lines(
62            config,
63            &mut (*self.eof_tok).whitespace_before.borrow_mut(),
64            Some(""),
65        )?;
66        let mut header = vec![];
67        if let Some(stmt) = body.first_mut() {
68            swap(stmt.leading_lines(), &mut header);
69            let mut last_indented = None;
70            for (num, line) in footer.iter().enumerate() {
71                if !line.whitespace.0.is_empty() {
72                    last_indented = Some(num);
73                } else if line.comment.is_some() {
74                    // This is a non-indented comment. Everything from here should belong in the
75                    // footer.
76                    break;
77                }
78            }
79            if let Some(num) = last_indented {
80                let (_, rest) = footer.split_at(num);
81                footer = rest.to_vec();
82            }
83        } else {
84            swap(&mut header, &mut footer);
85        }
86        Ok(Self::Inflated {
87            body,
88            header,
89            footer,
90            default_indent,
91            default_newline,
92            has_trailing_newline,
93            encoding: self.encoding,
94        })
95    }
96}