From 9f4527c10b43359f5a76d69b5f57391a686f92ee Mon Sep 17 00:00:00 2001
From: "Carson M." <carson@pyke.io>
Date: Wed, 28 Aug 2024 12:25:32 -0500
Subject: [PATCH] feat: get all custom metadata keys, closes #277

---
 src/metadata.rs     | 20 ++++++++++++++++++++
 tests/vectorizer.rs |  1 +
 2 files changed, 21 insertions(+)

diff --git a/src/metadata.rs b/src/metadata.rs
index 84fc69e2..9625b9df 100644
--- a/src/metadata.rs
+++ b/src/metadata.rs
@@ -92,6 +92,26 @@ impl<'s> ModelMetadata<'s> {
 			Ok(None)
 		}
 	}
+
+	pub fn custom_keys(&self) -> Result<Vec<String>> {
+		let mut keys: *mut *mut c_char = std::ptr::null_mut();
+		let mut key_len = 0;
+		ortsys![unsafe ModelMetadataGetCustomMetadataMapKeys(self.metadata_ptr.as_ptr(), self.allocator.ptr.as_ptr(), &mut keys, &mut key_len) -> Error::GetModelMetadata];
+		if key_len != 0 && !keys.is_null() {
+			let res = unsafe { std::slice::from_raw_parts(keys, key_len as usize) }
+				.iter()
+				.map(|c| {
+					let res = char_p_to_string(*c);
+					unsafe { self.allocator.free(*c) };
+					res
+				})
+				.collect();
+			unsafe { self.allocator.free(keys) };
+			res
+		} else {
+			Ok(vec![])
+		}
+	}
 }
 
 impl<'s> Drop for ModelMetadata<'s> {
diff --git a/tests/vectorizer.rs b/tests/vectorizer.rs
index f3af20a7..ed74d636 100644
--- a/tests/vectorizer.rs
+++ b/tests/vectorizer.rs
@@ -17,6 +17,7 @@ fn vectorizer() -> ort::Result<()> {
 	let metadata = session.metadata()?;
 	assert_eq!(metadata.producer()?, "skl2onnx");
 	assert_eq!(metadata.description()?, "test description");
+	assert_eq!(metadata.custom_keys()?, ["custom_key"]);
 	assert_eq!(metadata.custom("custom_key")?.as_deref(), Some("custom_value"));
 
 	let array = ndarray::CowArray::from(ndarray::Array::from_shape_vec((1,), vec!["document".to_owned()]).unwrap());