use std::fs::File; use std::io::{BufWriter, Write}; use std::path::Path; use erg_common::io::Input; use erg_common::pathutil::mod_name; use erg_common::set::Set; use erg_common::traits::LimitedDisplay; use erg_common::{log, Str}; use erg_compiler::build_package::{CheckStatus, PylyzerStatus}; use erg_compiler::hir::{ClassDef, Expr, HIR}; use erg_compiler::ty::value::{GenTypeObj, TypeObj}; use erg_compiler::ty::{HasType, Type}; pub struct DeclFile { pub filename: String, pub code: String, } fn escape_type(typ: String) -> String { typ.replace('%', "Type_").replace("", "") } pub struct DeclFileGenerator { filename: String, namespace: String, imported: Set, code: String, } impl DeclFileGenerator { pub fn new(input: &Input, status: CheckStatus) -> Self { let (timestamp, hash) = { let py_file_path = input.path(); let metadata = std::fs::metadata(py_file_path).unwrap(); let dummy_hash = metadata.len(); (metadata.modified().unwrap(), dummy_hash) }; let status = PylyzerStatus { status, file: input.path().into(), timestamp, hash, }; let code = format!("{status}\n"); Self { filename: input.filename().replace(".py", ".d.er"), namespace: "".to_string(), imported: Set::new(), code, } } pub fn gen_decl_er(mut self, hir: HIR) -> DeclFile { for chunk in hir.module.into_iter() { self.gen_chunk_decl(chunk); } log!("code:\n{}", self.code); DeclFile { filename: self.filename, code: self.code, } } fn gen_chunk_decl(&mut self, chunk: Expr) { match chunk { Expr::Def(def) => { let mut name = def .sig .ident() .inspect() .replace('\0', "") .replace(['%', '*'], "___"); let ref_t = def.sig.ident().ref_t(); let typ = escape_type(ref_t.replace_failure().to_string_unabbreviated()); // Erg can automatically import nested modules // `import http.client` => `http = pyimport "http"` let decl = if ref_t.is_py_module() { name = name.split('.').next().unwrap().to_string(); let full_path_str = ref_t.typarams()[0].to_string_unabbreviated(); let mod_name = mod_name(Path::new(full_path_str.trim_matches('"'))); let imported = if self.imported.insert(mod_name.clone()) { format!("{}.{mod_name} = pyimport \"{mod_name}\"", self.namespace) } else { "".to_string() }; if self.imported.insert(name.clone().into()) { format!( "{}.{name} = pyimport \"{mod_name}\"\n{imported}", self.namespace, ) } else { imported } } else { format!("{}.{name}: {typ}", self.namespace) }; self.code += &decl; } Expr::ClassDef(def) => { let class_name = def .sig .ident() .inspect() .replace('\0', "") .replace(['%', '*'], "___"); let src = format!("{}.{class_name}", self.namespace); let stash = std::mem::replace(&mut self.namespace, src); let decl = format!(".{class_name}: ClassType"); self.code += &decl; self.code.push('\n'); if let GenTypeObj::Subclass(class) = &def.obj { let sup = class .sup .as_ref() .typ() .replace_failure() .to_string_unabbreviated(); let sup = escape_type(sup); let decl = format!(".{class_name} <: {sup}\n"); self.code += &decl; } if let Some(TypeObj::Builtin { t: Type::Record(rec), .. }) = def.obj.base_or_sup() { for (attr, t) in rec.iter() { let typ = escape_type(t.replace_failure().to_string_unabbreviated()); let decl = format!("{}.{}: {typ}\n", self.namespace, attr.symbol); self.code += &decl; } } if let Some(TypeObj::Builtin { t: Type::Record(rec), .. }) = def.obj.additional() { for (attr, t) in rec.iter() { let typ = escape_type(t.replace_failure().to_string_unabbreviated()); let decl = format!("{}.{}: {typ}\n", self.namespace, attr.symbol); self.code += &decl; } } for attr in ClassDef::take_all_methods(def.methods_list) { self.gen_chunk_decl(attr); } self.namespace = stash; } Expr::Dummy(dummy) => { for chunk in dummy.into_iter() { self.gen_chunk_decl(chunk); } } _ => {} } self.code.push('\n'); } } pub fn reserve_decl_er(input: Input) { let mut dir = input.dir(); dir.push("__pycache__"); let pycache_dir = dir.as_path(); if !pycache_dir.exists() { std::fs::create_dir(pycache_dir).unwrap(); } let filename = input.filename(); let mut path = pycache_dir.join(filename); path.set_extension("d.er"); if !path.exists() { let _f = File::create(path).unwrap(); } } pub fn dump_decl_er(input: Input, hir: HIR, status: CheckStatus) { let decl_gen = DeclFileGenerator::new(&input, status); let file = decl_gen.gen_decl_er(hir); let mut dir = input.dir(); dir.push("__pycache__"); let pycache_dir = dir.as_path(); let f = File::options() .write(true) .open(pycache_dir.join(file.filename)) .unwrap(); let mut f = BufWriter::new(f); f.write_all(file.code.as_bytes()).unwrap(); }