Fix up tests

spiraldb · Mar 26, 2024 · 67a30c9 · 67a30c9
1 parent 3d7fe88
commit 67a30c9
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 17 deletions.
diff --git a/pyvortex/src/lib.rs b/pyvortex/src/lib.rs
@@ -3,6 +3,7 @@ use pyo3::prelude::*;
 
 use dtype::PyDType;
 use vortex_schema::DType;
+use vortex_schema::Signedness::{Signed, Unsigned};
 
 use crate::array::*;
 use crate::compress::PyCompressConfig;
@@ -56,6 +57,7 @@ fn _lib(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyCompressConfig>()?;
 
     m.add_function(wrap_pyfunction!(dtype_int, m)?)?;
+    m.add_function(wrap_pyfunction!(dtype_uint, m)?)?;
     m.add_function(wrap_pyfunction!(dtype_float, m)?)?;
     m.add_function(wrap_pyfunction!(dtype_bool, m)?)?;
     m.add_function(wrap_pyfunction!(dtype_utf8, m)?)?;
@@ -70,23 +72,30 @@ fn dtype_bool(py: Python<'_>, nullable: bool) -> PyResult<Py<PyDType>> {
 }
 
 #[pyfunction(name = "int")]
-#[pyo3(signature = (width = None, signed = true, nullable = false))]
-fn dtype_int(
-    py: Python<'_>,
-    width: Option<u16>,
-    signed: bool,
-    nullable: bool,
-) -> PyResult<Py<PyDType>> {
+#[pyo3(signature = (width = None, nullable = false))]
+fn dtype_int(py: Python<'_>, width: Option<u16>, nullable: bool) -> PyResult<Py<PyDType>> {
+    PyDType::wrap(
+        py,
+        DType::Int(width.unwrap_or(64).into(), Signed, nullable.into()),
+    )
+}
+
+#[pyfunction(name = "uint")]
+#[pyo3(signature = (width = None, nullable = false))]
+fn dtype_uint(py: Python<'_>, width: Option<u16>, nullable: bool) -> PyResult<Py<PyDType>> {
     PyDType::wrap(
         py,
-        DType::Int(width.unwrap_or(0).into(), signed.into(), nullable.into()),
+        DType::Int(width.unwrap_or(64).into(), Unsigned, nullable.into()),
     )
 }
 
 #[pyfunction(name = "float")]
 #[pyo3(signature = (width = None, nullable = false))]
 fn dtype_float(py: Python<'_>, width: Option<i8>, nullable: bool) -> PyResult<Py<PyDType>> {
-    PyDType::wrap(py, DType::Float(width.unwrap_or(0).into(), nullable.into()))
+    PyDType::wrap(
+        py,
+        DType::Float(width.unwrap_or(64).into(), nullable.into()),
+    )
 }
 
 #[pyfunction(name = "utf8")]

diff --git a/pyvortex/test/test_dtype.py b/pyvortex/test/test_dtype.py
@@ -2,8 +2,9 @@
 
 
 def test_int():
-    assert str(vortex.int()) == "signed_int(_)"
-    assert str(vortex.int(32)) == "signed_int(32)"
-    assert str(vortex.int(32, signed=False)) == "unsigned_int(32)"
+    assert str(vortex.int()) == "int(64)"
+    assert str(vortex.int(32)) == "int(32)"
+    assert str(vortex.int(32, nullable=True)) == "int(32)?"
+    assert str(vortex.uint(32)) == "uint(32)"
     assert str(vortex.float(16)) == "float(16)"
     assert str(vortex.bool(nullable=True)) == "bool?"
diff --git a/vortex-array/src/formatter.rs b/vortex-array/src/formatter.rs
@@ -120,15 +120,12 @@ mod test {
     #[test]
     fn display_primitive() {
         let arr: ArrayRef = (0..100).collect::<Vec<i32>>().into_array();
-        assert_eq!(
-            format!("{}", arr),
-            "vortex.primitive(signed_int(32), len=100)"
-        );
+        assert_eq!(format!("{}", arr), "vortex.primitive(int(32), len=100)");
     }
 
     #[test]
     fn tree_display_primitive() {
         let arr: ArrayRef = (0..100).collect::<Vec<i32>>().into_array();
-        assert_eq!(display_tree(&arr), "root: vortex.primitive(signed_int(32), len=100) nbytes=400 B (100.00%)\n  values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...\n  validity: None\n")
+        assert_eq!(display_tree(&arr), "root: vortex.primitive(int(32), len=100) nbytes=400 B (100.00%)\n  values: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...\n  validity: None\n")
     }
 }