rust/tools/gen_assigned_numbers.rs

// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! This tool generates Rust code with assigned number tables from the equivalent Python.

use pyo3::{
    intern,
    types::{PyDict, PyModule},
    PyResult, Python,
};
use std::{collections, env, fs, path};

fn main() -> anyhow::Result<()> {
    pyo3::prepare_freethreaded_python();
    let mut dir = path::Path::new(&env::var("CARGO_MANIFEST_DIR")?).to_path_buf();
    dir.push("src/wrapper/assigned_numbers");

    company_ids(&dir)?;

    Ok(())
}

fn company_ids(base_dir: &path::Path) -> anyhow::Result<()> {
    let mut sorted_ids = load_company_ids()?.into_iter().collect::<Vec<_>>();
    sorted_ids.sort_by_key(|(id, _name)| *id);

    let mut contents = String::new();
    contents.push_str(LICENSE_HEADER);
    contents.push_str("\n\n");
    contents.push_str(
        "// auto-generated by gen_assigned_numbers, do not edit

use crate::wrapper::core::Uuid16;
use lazy_static::lazy_static;
use std::collections;

lazy_static! {
    /// Assigned company IDs
    pub static ref COMPANY_IDS: collections::HashMap<Uuid16, &'static str> = [
",
    );

    for (id, name) in sorted_ids {
        contents.push_str(&format!("        ({id}_u16, r#\"{name}\"#),\n"))
    }

    contents.push_str(
        "    ]
    .into_iter()
    .map(|(id, name)| (Uuid16::from_be_bytes(id.to_be_bytes()), name))
    .collect();
}
",
    );

    let mut company_ids = base_dir.to_path_buf();
    company_ids.push("company_ids.rs");
    fs::write(&company_ids, contents)?;

    Ok(())
}

fn load_company_ids() -> PyResult<collections::HashMap<u16, String>> {
    Python::with_gil(|py| {
        PyModule::import(py, intern!(py, "bumble.company_ids"))?
            .getattr(intern!(py, "COMPANY_IDENTIFIERS"))?
            .downcast::<PyDict>()?
            .into_iter()
            .map(|(k, v)| Ok((k.extract::<u16>()?, v.str()?.to_str()?.to_string())))
            .collect::<PyResult<collections::HashMap<_, _>>>()
    })
}

const LICENSE_HEADER: &str = r#"// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License."#;