feat: add snowflake

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fastid"
-version = "0.0.1"
+version = "0.0.2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -12,4 +12,6 @@ crate-type = ["cdylib"]
 
 [dependencies]
 ulid="1.1.0"
-pyo3="0.20.0"
+pyo3="0.20.0"
+rs-snowflake = "0.6.0"
+once_cell = "1.19.0"

fastid.pyi

@@ -1,2 +1,8 @@
 def ulid() -> str:
-    ...
+    ...
+
+def snowflake_str(machine_id: int = 1, node_id: int = 1) -> str:
+    ...
+
+def snowflake_int(machine_id: int = 1, node_id: int = 1) -> int:
+    ...

python/test/test_fastid.py

@@ -6,5 +6,11 @@
 
 def test_ulid():
     print()
-    print(timeit.timeit(lambda: microsecond.new().str, number=10000))
-    print(timeit.timeit(fastid.ulid, number=10000))
+    print("ulid-py    ", timeit.timeit(lambda: microsecond.new().str, number=10000))
+    print("fastid.ulid", timeit.timeit(lambda: fastid.ulid(), number=10000))
+
+
+def test_snowflake():
+    print()
+    print("fastid.snowflake_int", timeit.timeit(lambda: fastid.snowflake_int(), number=10000))
+    print("fastid.snowflake_str", timeit.timeit(lambda: fastid.snowflake_str(), number=10000))

src/lib.rs

@@ -1,14 +1,44 @@
+use std::sync::Mutex;
+use once_cell::sync::Lazy;
 use pyo3::prelude::*;
-use ulid::{Ulid};
+use snowflake::SnowflakeIdGenerator;
+use ulid::Ulid;
 
+static SNOWFLAKE_GENERATOR: Lazy<Mutex<SnowflakeIdGenerator>> = Lazy::new(|| {
+    Mutex::new(SnowflakeIdGenerator::new(1, 1))
+});
 #[pyfunction]
-#[pyo3(name="ulid")]
+#[pyo3(name = "ulid")]
 fn get_ulid() -> PyResult<String> {
     Ok(Ulid::new().to_string())
 }
 
+
+fn get_snowflake(machine_id: Option<i32>, node_id: Option<i32>) -> i64 {
+    let mut sg = SNOWFLAKE_GENERATOR.lock().unwrap();
+    sg.node_id = node_id.unwrap_or(1);
+    sg.machine_id = machine_id.unwrap_or(1);
+    sg.real_time_generate()
+}
+#[pyfunction]
+#[pyo3(name = "snowflake_str")]
+fn get_snowflake_str(machine_id: Option<i32>, node_id: Option<i32>) -> PyResult<String> {
+    Ok(format!("{:X}", get_snowflake(machine_id, node_id)))
+}
+
+#[pyfunction]
+#[pyo3(name = "snowflake_int")]
+fn get_snowflake_int(machine_id: Option<i32>, node_id: Option<i32>) -> PyResult<i64> {
+    Ok(get_snowflake(machine_id, node_id))
+}
+
+
 #[pymodule]
 fn fastid(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(get_ulid, m)?)?;
+
+    m.add_function(wrap_pyfunction!(get_snowflake_int, m)?)?;
+    m.add_function(wrap_pyfunction!(get_snowflake_str, m)?)?;
+
     Ok(())
 }