add version getter function

Author: blackerby

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "polars-capitol"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 
 [lib]
@@ -12,5 +12,5 @@ pyo3 = { version = "0.22", features = ["extension-module", "abi3-py38"] }
 pyo3-polars = { version = "0.19.0", features = ["derive"] }
 serde = { version = "1", features = ["derive"] }
 polars = { version = "0.45.1", default-features = false }
-capitol = "0.2.0"
+capitol = "0.2.1"
 

Cargo.toml

@@ -21,3 +21,12 @@ def cdg_url(expr: IntoExprColumn) -> pl.Expr:
         function_name="cdg_url",
         is_elementwise=True,
     )
+
+
+def version(expr: IntoExprColumn) -> pl.Expr:
+    return register_plugin_function(
+        args=[expr],
+        plugin_path=LIB,
+        function_name="version",
+        is_elementwise=True,
+    )

polars_capitol/__init__.py

@@ -15,5 +15,7 @@
 s = pl.col("package_id").str.split(by="-").list.get(1, null_on_oob=True)
 expr = cap.cdg_url(s)
 result = df.with_columns(cdg_url=expr)
+expr = cap.version(s)
+result = df.with_columns(version=expr)
 
 print(result)

run.py

@@ -14,3 +14,21 @@ fn cdg_url(inputs: &[Series]) -> PolarsResult<Series> {
     });
     Ok(out.into_series())
 }
+
+#[polars_expr(output_type=String)]
+fn version(inputs: &[Series]) -> PolarsResult<Series> {
+    let s = &inputs[0];
+    let ca = s.str()?;
+    let out = ca.try_apply_into_string_amortized(|value: &str, output: &mut String| {
+        let citation = Citation::parse(value).unwrap();
+        if let Some(version) = citation.version() {
+            write!(output, "{}", version).unwrap();
+            Ok(())
+        } else {
+            Err(PolarsError::InvalidOperation(
+                "`version` called on citation without version".into(),
+            ))
+        }
+    });
+    Ok(out?.into_series())
+}

src/expressions.rs

@@ -0,0 +1,49 @@
+import polars as pl
+import polars_capitol as cap
+import pytest
+
+
+def test_version():
+    df = pl.DataFrame(
+        {
+            "package_id": [
+                "BILLS-118hr8070ih",
+            ]
+        }
+    )
+    s = pl.col("package_id").str.split(by="-").list.get(1, null_on_oob=True)
+    expr = cap.version(s)
+    result = df.with_columns(version=expr)
+
+    expected_df = pl.DataFrame(
+        {
+            "package_id": [
+                "BILLS-118hr8070ih",
+            ],
+            "version": [
+                "ih",
+            ],
+        }
+    )
+
+    assert result.equals(expected_df)
+
+
+def test_citation_without_version_raises():
+    df = pl.DataFrame(
+        {
+            "package_id": [
+                "BILLS-118hr8070",
+            ]
+        }
+    )
+    s = pl.col("package_id").str.split(by="-").list.get(1, null_on_oob=True)
+    expr = cap.version(s)
+
+    with pytest.raises(pl.exceptions.ComputeError) as exception:
+        df.with_columns(version=expr)
+
+    assert (
+        str(exception.value)
+        == "the plugin failed with message: `version` called on citation without version"
+    )