Skip to main content

probl_engine/data/
schema.rs

1//! `probl schema FILE`: a type to read a file with, guessed from what's in
2//! it (docs/data-input.md). It's a suggestion for the program's author to
3//! paste and edit: running a program never guesses. It reads the file
4//! within the same limits as loading it.
5
6use super::InputLimits;
7use crate::dates;
8use crate::report::thousands;
9use csv::{ReaderBuilder, StringRecord, Trim};
10use probl_sema::data::field_key;
11use probl_sema::ir::DataFormat;
12use probl_syntax::token::is_name;
13
14/// Suggest a declaration, and the `read` that uses it, for the data in
15/// `bytes`, read from `file` in `format`.
16pub fn suggest(bytes: &[u8], format: DataFormat, file: &str, limits: &InputLimits) -> Result<String, String> {
17    if bytes.len() as u64 > limits.max_bytes {
18        return Err(format!(
19            "the file is more than {} MiB",
20            limits.max_bytes.div_ceil(1024 * 1024)
21        ));
22    }
23    let bytes = bytes.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(bytes);
24    let stem = file.rsplit(['/', '\\']).next().unwrap_or(file);
25    let stem = stem.rsplit_once('.').map_or(stem, |(s, _)| s);
26    let fallback = if format == DataFormat::Csv { "Row" } else { "Data" };
27    let type_name = camel(stem).unwrap_or_else(|| fallback.to_string());
28    let var = field_name(stem).unwrap_or_else(|| "data".to_string());
29    let mut out = Output {
30        notes: Vec::new(),
31        values: 0,
32        limits,
33    };
34    let text = match format {
35        DataFormat::Csv => {
36            let fields = csv_fields(bytes, &mut out)?;
37            format!(
38                "type {type_name} = {{ {} }}\n# let {var}: list[{type_name}] = read({file:?})\n",
39                fields.join(", ")
40            )
41        }
42        DataFormat::Lines => {
43            let text =
44                std::str::from_utf8(bytes).map_err(|_| "the file isn't text: it isn't valid UTF-8".to_string())?;
45            let mut column = Column::default();
46            for line in text.lines().filter(|l| !l.trim().is_empty()) {
47                out.count()?;
48                column.add(line, out.limits.max_integer_bits);
49            }
50            format!("# let {var}: list[{}] = read({file:?})\n", column.guess())
51        }
52        DataFormat::Json => {
53            let value: serde_json::Value =
54                serde_json::from_slice(bytes).map_err(|e| format!("this isn't valid JSON: {e}"))?;
55            let shape = shape(&value, "", &mut out)?;
56            match &shape {
57                Shape::Record(_) => format!(
58                    "type {type_name} = {}\n# let {var}: {type_name} = read({file:?})\n",
59                    render(&shape, &mut out, "")
60                ),
61                Shape::List(item) if matches!(**item, Shape::Record(_)) => format!(
62                    "type {type_name} = {}\n# let {var}: list[{type_name}] = read({file:?})\n",
63                    render(item, &mut out, "")
64                ),
65                _ => format!("# let {var}: {} = read({file:?})\n", render(&shape, &mut out, "")),
66            }
67        }
68    };
69    let mut notes = out.notes;
70    notes.dedup();
71    Ok(notes.iter().map(|n| format!("# {n}\n")).collect::<String>() + &text)
72}
73
74/// Notes for the reader, and the values counted against the limit.
75struct Output<'a> {
76    notes: Vec<String>,
77    values: u64,
78    limits: &'a InputLimits,
79}
80
81impl Output<'_> {
82    fn count(&mut self) -> Result<(), String> {
83        self.values += 1;
84        if self.values > self.limits.max_values {
85            return Err(format!(
86                "the file has more than {} values",
87                thousands(self.limits.max_values.min(i64::MAX as u64) as i64)
88            ));
89        }
90        Ok(())
91    }
92}
93
94/// A field name for a column or key, matching it, if there is one.
95fn field_name(name: &str) -> Option<String> {
96    let mut words: Vec<String> = Vec::new();
97    let mut word = String::new();
98    let mut previous_lower = false;
99    for c in name.chars() {
100        if c.is_alphanumeric() {
101            if c.is_uppercase() && previous_lower && !word.is_empty() {
102                words.push(std::mem::take(&mut word));
103            }
104            previous_lower = c.is_lowercase() || c.is_ascii_digit();
105            word.extend(c.to_lowercase());
106        } else {
107            if !word.is_empty() {
108                words.push(std::mem::take(&mut word));
109            }
110            previous_lower = false;
111        }
112    }
113    if !word.is_empty() {
114        words.push(word);
115    }
116    let snake = words.join("_");
117    [snake.clone(), format!("{snake}_")]
118        .into_iter()
119        .find(|n| is_name(n) && field_key(n) == field_key(name))
120}
121
122/// A type name for a file name: `sales-2026` gives `Sales2026`.
123fn camel(stem: &str) -> Option<String> {
124    let name: String = stem
125        .split(|c: char| !c.is_ascii_alphanumeric())
126        .filter(|w| !w.is_empty())
127        .map(|w| {
128            let mut cs = w.chars();
129            let first = cs.next().unwrap().to_ascii_uppercase();
130            std::iter::once(first).chain(cs).collect::<String>()
131        })
132        .collect();
133    is_name(&name).then_some(name)
134}
135
136/// What a column of text could be: the first of these types that reads
137/// every value, ignoring empty ones.
138#[derive(Clone, Copy, Debug)]
139struct Column {
140    int: bool,
141    float: bool,
142    prob: bool,
143    bool_: bool,
144    date: bool,
145    values: bool,
146    empty: bool,
147}
148
149impl Default for Column {
150    fn default() -> Column {
151        Column {
152            int: true,
153            float: true,
154            prob: true,
155            bool_: true,
156            date: true,
157            values: false,
158            empty: false,
159        }
160    }
161}
162
163impl Column {
164    fn add(&mut self, text: &str, max_integer_bits: u64) {
165        let t = text.trim();
166        if t.is_empty() {
167            self.empty = true;
168            return;
169        }
170        self.values = true;
171        self.int &= t
172            .parse::<probl_number::Integer>()
173            .is_ok_and(|n| n.bits() <= max_integer_bits);
174        self.float &= t.parse::<f64>().is_ok_and(f64::is_finite);
175        self.prob &= t
176            .strip_suffix('%')
177            .and_then(|p| p.trim().parse::<f64>().ok())
178            .is_some_and(|x| (0.0..=100.0).contains(&x));
179        self.bool_ &= t.eq_ignore_ascii_case("true") || t.eq_ignore_ascii_case("false");
180        self.date &= dates::parse(t).is_some();
181    }
182
183    fn guess(&self) -> &'static str {
184        match self {
185            _ if !self.values => "str",
186            Column { int: true, .. } => "int",
187            Column { float: true, .. } => "float",
188            Column { prob: true, .. } => "prob",
189            Column { bool_: true, .. } => "bool",
190            Column { date: true, .. } => "date",
191            _ => "str",
192        }
193    }
194}
195
196fn csv_fields(bytes: &[u8], out: &mut Output) -> Result<Vec<String>, String> {
197    let mut reader = ReaderBuilder::new()
198        .has_headers(true)
199        .flexible(false)
200        .trim(Trim::None)
201        .from_reader(bytes);
202    let headers = reader
203        .headers()
204        .map_err(|e| format!("can't read it as CSV: {e}"))?
205        .clone();
206    if headers.is_empty() {
207        return Err("the file is empty: a CSV file starts with a row of column names".to_string());
208    }
209    let mut columns = vec![Column::default(); headers.len()];
210    let mut record = StringRecord::new();
211    while reader
212        .read_record(&mut record)
213        .map_err(|e| format!("can't read it as CSV: {e}"))?
214    {
215        for (column, cell) in columns.iter_mut().zip(record.iter()) {
216            out.count()?;
217            column.add(cell, out.limits.max_integer_bits);
218        }
219    }
220    let mut fields = Vec::new();
221    let mut keys: Vec<(String, String)> = Vec::new();
222    for (header, column) in headers.iter().zip(&columns) {
223        let Some(name) = field_name(header) else {
224            out.notes.push(format!(
225                "the column {header:?} can't be a field name: rename it in the file to read it"
226            ));
227            continue;
228        };
229        if let Some((other, _)) = keys.iter().find(|(_, k)| *k == field_key(header)) {
230            out.notes.push(format!(
231                "the columns {other:?} and {header:?} would match the same field: rename one to read either"
232            ));
233            continue;
234        }
235        keys.push((header.to_string(), field_key(header)));
236        let ty = column.guess();
237        if column.empty && ty != "str" {
238            out.notes.push(format!(
239                "`{name}` has empty cells, which only a `str` field reads for now"
240            ));
241        }
242        fields.push(format!("{name}: {ty}"));
243    }
244    Ok(fields)
245}
246
247/// The type a JSON value suggests.
248#[derive(Clone, Debug, PartialEq)]
249enum Shape {
250    /// Nothing to go on: an empty array, or only nulls.
251    Unknown,
252    Int,
253    Float,
254    Prob,
255    Bool,
256    Date,
257    Str,
258    List(Box<Shape>),
259    Map(&'static str, Box<Shape>),
260    Record(Vec<(String, Shape)>),
261    /// Values of different kinds.
262    Mixed,
263}
264
265fn shape(value: &serde_json::Value, at: &str, out: &mut Output) -> Result<Shape, String> {
266    use serde_json::Value as J;
267    out.count()?;
268    Ok(match value {
269        J::Null => {
270            out.notes.push(format!(
271                "{} has nulls: the language has no missing values yet",
272                place(at)
273            ));
274            Shape::Unknown
275        }
276        J::Bool(_) => Shape::Bool,
277        J::Number(n) => {
278            let text = n.to_string();
279            if !text.contains(['.', 'e', 'E']) {
280                let n = text.parse::<probl_number::Integer>().map_err(|e| e.to_string())?;
281                if n.bits() > out.limits.max_integer_bits {
282                    return Err("integer size exceeds the input limit".into());
283                }
284                Shape::Int
285            } else {
286                Shape::Float
287            }
288        }
289        J::String(s) => {
290            let mut column = Column::default();
291            column.add(s, out.limits.max_integer_bits);
292            match column.guess() {
293                "date" => Shape::Date,
294                "prob" => Shape::Prob,
295                _ => Shape::Str,
296            }
297        }
298        J::Array(items) => {
299            let mut item = Shape::Unknown;
300            for (i, x) in items.iter().enumerate() {
301                item = unify(item, shape(x, &format!("{at}[{i}]"), out)?);
302            }
303            Shape::List(Box::new(item))
304        }
305        J::Object(entries) => {
306            // Keys that are all dates or all integers make a map.
307            let key_type = if entries.is_empty() {
308                None
309            } else if entries.keys().all(|k| dates::parse(k).is_some()) {
310                Some("date")
311            } else if entries.keys().all(|k| {
312                k.parse::<probl_number::Integer>()
313                    .is_ok_and(|n| n.bits() <= out.limits.max_integer_bits)
314            }) {
315                Some("int")
316            } else {
317                None
318            };
319            match key_type {
320                Some(k) => {
321                    let mut value = Shape::Unknown;
322                    for (key, x) in entries {
323                        value = unify(value, shape(x, &format!("{at}.{key}"), out)?);
324                    }
325                    Shape::Map(k, Box::new(value))
326                }
327                None => {
328                    let mut fields = Vec::new();
329                    for (key, x) in entries {
330                        let path = if at.is_empty() {
331                            key.clone()
332                        } else {
333                            format!("{at}.{key}")
334                        };
335                        fields.push((key.clone(), shape(x, &path, out)?));
336                    }
337                    Shape::Record(fields)
338                }
339            }
340        }
341    })
342}
343
344fn place(at: &str) -> String {
345    if at.is_empty() {
346        "the document".to_string()
347    } else {
348        format!("`{at}`")
349    }
350}
351
352fn unify(a: Shape, b: Shape) -> Shape {
353    use Shape::*;
354    match (a, b) {
355        (Unknown, x) | (x, Unknown) => x,
356        (a, b) if a == b => a,
357        (Int, Float) | (Float, Int) => Float,
358        (Date | Prob | Str, Date | Prob | Str) => Str,
359        (List(a), List(b)) => List(Box::new(unify(*a, *b))),
360        (Map(k, a), Map(l, b)) if k == l => Map(k, Box::new(unify(*a, *b))),
361        (Record(a), Record(b)) => {
362            let mut fields = a;
363            for (name, shape) in b {
364                match fields.iter_mut().find(|(n, _)| *n == name) {
365                    Some((_, s)) => *s = unify(std::mem::replace(s, Unknown), shape),
366                    None => fields.push((name, shape)),
367                }
368            }
369            Record(fields)
370        }
371        _ => Mixed,
372    }
373}
374
375fn render(shape: &Shape, out: &mut Output, at: &str) -> String {
376    match shape {
377        Shape::Unknown => {
378            out.notes
379                .push(format!("{} gives nothing to guess from: `str` is a guess", place(at)));
380            "str".into()
381        }
382        Shape::Int => "int".into(),
383        Shape::Float => "float".into(),
384        Shape::Prob => "prob".into(),
385        Shape::Bool => "bool".into(),
386        Shape::Date => "date".into(),
387        Shape::Str => "str".into(),
388        Shape::Mixed => {
389            out.notes.push(format!(
390                "{} mixes different kinds of values: `str` is a guess",
391                place(at)
392            ));
393            "str".into()
394        }
395        Shape::List(item) => format!("list[{}]", render(item, out, &format!("{at}[]"))),
396        Shape::Map(k, v) => format!("map[{k}, {}]", render(v, out, &format!("{at}.*"))),
397        Shape::Record(fields) => render_fields(fields, out, at),
398    }
399}
400
401fn render_fields(fields: &[(String, Shape)], out: &mut Output, at: &str) -> String {
402    let mut parts = Vec::new();
403    let mut keys: Vec<String> = Vec::new();
404    for (key, shape) in fields {
405        let path = if at.is_empty() {
406            key.clone()
407        } else {
408            format!("{at}.{key}")
409        };
410        let Some(name) = field_name(key) else {
411            out.notes
412                .push(format!("the key {key:?} can't be a field name: it's left out"));
413            continue;
414        };
415        if keys.contains(&field_key(key)) {
416            out.notes.push(format!(
417                "the key {key:?} would match the same field as another: it's left out"
418            ));
419            continue;
420        }
421        keys.push(field_key(key));
422        parts.push(format!("{name}: {}", render(shape, out, &path)));
423    }
424    format!("{{ {} }}", parts.join(", "))
425}