Almost done WIP of exposing python types and docstrings

bin/stub_gen.rs

@@ -0,0 +1,223 @@
+extern crate mwalib;
+use std::env;
+use std::fs::File;
+use std::io::{Read, Write};
+
+fn main() -> anyhow::Result<()> {
+    env_logger::Builder::from_env(env_logger::Env::default().filter_or("RUST_LOG", "info")).init();
+
+    // Generating the stub requires the below env variable to be set for some reason?
+    env::set_var("CARGO_MANIFEST_DIR", env::current_dir()?);
+    let stub = mwalib::python::stub_info()?;
+    stub.generate()?;
+
+    // After the stub is generated, we have some "manual" fixes to do
+    let stubfile = String::from("mwalib.pyi");
+
+    // Import datetime
+    insert_stub_below(
+        &stubfile,
+        "import typing\n",
+        "from datetime import datetime\n",
+    )?;
+
+    // Add sched_start_utc (Chrono::DateTime<FixedOffset> is not supported yet)
+    insert_stub_below(
+        &stubfile,
+        "    sched_end_unix_time_ms: int\n",
+        "    sched_start_utc: datetime.datetime\n",
+    )?;
+
+    // Add sched_end_utc (Chrono::DateTime<FixedOffset> is not supported yet)
+    insert_stub_below(
+        &stubfile,
+        "    sched_start_utc: datetime.datetime\n",
+        "    sched_end_utc: datetime.datetime\n",
+    )?;
+
+    // Replace the constructors as pyo3_stub_gen seems to ignore the text_signature
+    replace_stub(
+        &stubfile,
+        "def __new__(cls,metafits_filename,mwa_version = ...): ...",
+        "def __new__(cls, metafits_filename: str, mwa_version: typing.Optional[MWAVersion]=None)->MetafitsContext:\n        ...\n",
+    )?;
+
+    replace_stub(&stubfile, "def __new__(cls,metafits_filename,gpubox_filenames): ...", "def __new__(cls, metafits_filename: str, gpubox_filenames: list[str])->CorrelatorContext:\n        ...\n")?;
+
+    replace_stub(
+        &stubfile,
+        "def __new__(cls,metafits_filename,voltage_filenames): ...",
+        "def __new__(cls, metafits_filename: str, voltage_filenames: list[str])->VoltageContext:\n        ...\n",
+    )?;
+
+    replace_stub(
+        &stubfile,
+        "def __enter__(self, slf:MetafitsContext) -> MetafitsContext:",
+        "def __enter__(self) -> MetafitsContext:",
+    )?;
+
+    replace_stub(
+        &stubfile,
+        "def __enter__(self, slf:CorrelatorContext) -> CorrelatorContext:",
+        "def __enter__(self) -> CorrelatorContext:",
+    )?;
+
+    replace_stub(
+        &stubfile,
+        "def __enter__(self, slf:VoltageContext) -> VoltageContext:",
+        "def __enter__(self) -> VoltageContext:",
+    )?;
+
+    anyhow::Ok(())
+}
+
+/// Inserts new items in the stubfile for when pyo3 cant do it
+/// properly.
+///
+/// This will:
+/// * Open the created stubfile
+/// * Find the line in `string_to_find` (fails if not found)
+/// * Add a newline and `string_to_add_below` (fails if cannot)
+///
+/// # Arguments
+///
+/// * `stubfile` - filename of the stubfile to edit
+///
+/// * `string_to_find` - string to find, so we know where in the file to make the insert
+///
+/// * `string_to_add_below` - string to add on a new line below `string_to_find`
+///
+///
+/// # Returns
+///
+/// * Result Ok if stub file was modified successfully.
+///
+///
+fn insert_stub_below(
+    stubfile: &str,
+    string_to_find: &str,
+    string_to_add_below: &str,
+) -> anyhow::Result<()> {
+    // Open and read the file entirely
+    let mut src = File::open(stubfile)?;
+    let mut data = String::new();
+    src.read_to_string(&mut data)?;
+    drop(src); // Close the file early
+
+    // Run the replace operation in memory
+    let mut new_string: String = string_to_find.to_owned();
+    new_string.push_str(string_to_add_below);
+    let new_data = data.replace(string_to_find, &new_string);
+
+    // Recreate the file and dump the processed contents to it
+    let mut dst = File::create(stubfile)?;
+    dst.write_all(new_data.as_bytes())?;
+
+    anyhow::Ok(())
+}
+
+/// Replaces items in the stubfile for when pyo3 cant do it
+/// properly.
+///
+/// This will:
+/// * Open the created stubfile
+/// * Find the line in `string_to_find` (fails if not found)
+/// * Replace it with `string_to_replace` (fails if cannot)
+///
+/// # Arguments
+///
+/// * `stubfile` - filename of the stubfile to edit
+///
+/// * `string_to_find` - string to find (which will be replaced)
+///
+/// * `string_to_replace` - string to put in `string_to_find`s place
+///
+///
+/// # Returns
+///
+/// * Result Ok if stub file was modified successfully.
+///
+///
+fn replace_stub(
+    stubfile: &str,
+    string_to_find: &str,
+    string_to_replace: &str,
+) -> anyhow::Result<()> {
+    // Open and read the file entirely
+    let mut src = File::open(stubfile)?;
+    let mut data = String::new();
+    src.read_to_string(&mut data)?;
+    drop(src); // Close the file early
+
+    // Run the replace operation in memory
+    let new_data = data.replace(string_to_find, string_to_replace);
+
+    // Recreate the file and dump the processed contents to it
+    let mut dst = File::create(stubfile)?;
+    dst.write_all(new_data.as_bytes())?;
+
+    anyhow::Ok(())
+}
+
+#[test]
+fn test_insert_stub_below() {
+    // Create a test file
+    let text_filename: String = "test.pyi".to_string();
+
+    let mut test_file =
+        File::create(&text_filename).unwrap_or_else(|_| panic!("Could not open {}", text_filename));
+    // Write some lines
+    let content = "    Hello\n    World\n    1234";
+    let _ = test_file.write_all(content.as_bytes());
+
+    // Now run the add_stub command
+    insert_stub_below(&text_filename, "    World\n", "    added_string\n")
+        .unwrap_or_else(|_| panic!("Could not add_stub to {}", text_filename));
+
+    // Now reread the file
+    let test_file =
+        File::open(&text_filename).unwrap_or_else(|_| panic!("Could not open {}", text_filename));
+    let mut lines = Vec::new();
+
+    for line in std::io::read_to_string(test_file).unwrap().lines() {
+        lines.push(line.to_string())
+    }
+
+    assert_eq!(lines[0], "    Hello");
+    assert_eq!(lines[1], "    World");
+    assert_eq!(lines[2], "    added_string");
+    assert_eq!(lines[3], "    1234");
+}
+
+#[test]
+fn test_replace_stub() {
+    // Create a test file
+    let text_filename: String = "test.pyi".to_string();
+
+    let mut test_file =
+        File::create(&text_filename).unwrap_or_else(|_| panic!("Could not open {}", text_filename));
+    // Write some lines
+    let content = "    Hello\n    World\n    1234";
+    let _ = test_file.write_all(content.as_bytes());
+
+    // Now run the add_stub command
+    replace_stub(
+        &text_filename,
+        "    World\n",
+        "    This is the replaced string\n",
+    )
+    .unwrap_or_else(|_| panic!("Could not add_stub to {}", text_filename));
+
+    // Now reread the file
+    let test_file =
+        File::open(&text_filename).unwrap_or_else(|_| panic!("Could not open {}", text_filename));
+    let mut lines = Vec::new();
+
+    for line in std::io::read_to_string(test_file).unwrap().lines() {
+        lines.push(line.to_string())
+    }
+
+    assert_eq!(lines[0], "    Hello");
+    assert_eq!(lines[1], "    This is the replaced string");
+    assert_eq!(lines[2], "    1234");
+}