use std::fs::{create_dir_all, File}; use std::io::{BufWriter, Write}; use std::path::Path; use erg_common::pathutil::{mod_name, NormalizedPathBuf}; 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::module::SharedModuleCache; use erg_compiler::ty::value::{GenTypeObj, TypeObj}; use erg_compiler::ty::{HasType, Type}; pub struct DeclFile { pub filename: String, pub code: String, } pub struct DeclFileGenerator { filename: String, namespace: String, imported: Set, code: String, } impl DeclFileGenerator { pub fn new(path: &NormalizedPathBuf, status: CheckStatus) -> Self { let (timestamp, hash) = { let metadata = std::fs::metadata(path).unwrap(); let dummy_hash = metadata.len(); (metadata.modified().unwrap(), dummy_hash) }; let status = PylyzerStatus { status, file: path.to_path_buf(), timestamp, hash, }; let code = format!("{status}\n"); Self { filename: path .file_name() .unwrap() .to_string_lossy() .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.iter() { self.gen_chunk_decl(chunk); } log!("code:\n{}", self.code); DeclFile { filename: self.filename, code: self.code, } } fn escape_type(&self, typ: String) -> String { typ.replace('%', "Type_") .replace("", "") .trim_start_matches(self.filename.trim_end_matches(".d.er")) .trim_start_matches(&self.namespace) .to_string() } // e.g. `x: foo.Bar` => `foo = pyimport "foo"; x: foo.Bar` fn prepare_using_type(&mut self, typ: &Type) { let namespace = Str::rc(typ.namespace().split('.').next().unwrap()); if namespace != self.namespace && !namespace.is_empty() && self.imported.insert(namespace.clone()) { self.code += &format!("{namespace} = pyimport \"{namespace}\"\n"); } } 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(); self.prepare_using_type(ref_t); let typ = self.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(); self.prepare_using_type(class.sup.typ()); let sup = self.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() { self.prepare_using_type(t); let typ = self.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() { self.prepare_using_type(t); let typ = self.escape_type(t.replace_failure().to_string_unabbreviated()); let decl = format!("{}.{}: {typ}\n", self.namespace, attr.symbol); self.code += &decl; } } for attr in ClassDef::get_all_methods(&def.methods_list) { self.gen_chunk_decl(attr); } self.namespace = stash; } Expr::Dummy(dummy) => { for chunk in dummy.iter() { self.gen_chunk_decl(chunk); } } _ => {} } self.code.push('\n'); } } fn dump_decl_er(path: &NormalizedPathBuf, hir: &HIR, status: CheckStatus) { let decl_gen = DeclFileGenerator::new(path, status); let file = decl_gen.gen_decl_er(hir); let Some(dir) = path.parent().and_then(|p| p.canonicalize().ok()) else { return; }; let cache_dir = dir.join("__pycache__"); if !cache_dir.exists() { let _ = create_dir_all(&cache_dir); } let path = cache_dir.join(file.filename); if !path.exists() { let _f = File::create(&path); } let Ok(f) = File::options().write(true).open(path) else { return; }; let mut f = BufWriter::new(f); let _ = f.write_all(file.code.as_bytes()); } pub fn dump_decl_package(modules: &SharedModuleCache) { for (path, module) in modules.raw_iter() { if let Some(hir) = module.hir.as_ref() { dump_decl_er(path, hir, module.status); } } }