From edfe18b4870925bd582b800f30e1977229042f62 Mon Sep 17 00:00:00 2001
From: Abhijit Gadgil <gabhijit@iitbombay.org>
Date: Wed, 26 Jun 2024 15:16:47 +0530
Subject: [PATCH] Using tokio style cfg_* macros

Defined two macros `cfg_python!` and `cfg_wasm!` to wrap items within
feature specific blocks. Looks a lot cleaner

Signed-off-by: Abhijit Gadgil <gabhijit@iitbombay.org>
---
 src/cfg_macros.rs | 19 +++++++++++++++++++
 src/errors.rs     | 11 ++++++-----
 src/lib.rs        | 41 ++++++++++++++++++++++-------------------
 src/packet.rs     | 42 +++++++++++++++++++++---------------------
 4 files changed, 68 insertions(+), 45 deletions(-)
 create mode 100644 src/cfg_macros.rs

diff --git a/src/cfg_macros.rs b/src/cfg_macros.rs
new file mode 100644
index 0000000..e9b9ec1
--- /dev/null
+++ b/src/cfg_macros.rs
@@ -0,0 +1,19 @@
+#![allow(unused_macros)]
+
+macro_rules! cfg_python {
+    ($($item:item)*) => {
+        $(
+            #[cfg(feature = "python-bindings")]
+            $item
+        )*
+    }
+}
+
+macro_rules! cfg_wasm {
+    ($($item:item)*) => {
+        $(
+            #[cfg(feature = "wasm")]
+            $item
+        )*
+    }
+}
diff --git a/src/errors.rs b/src/errors.rs
index b876cc8..85baa14 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -30,10 +30,11 @@ impl core::fmt::Display for Error {
 }
 
 // Python Bindings
-#[cfg(feature = "python-bindings")]
-impl std::convert::From<Error> for pyo3::PyErr {
-    // TODO: Add proper error reporting
-    fn from(_e: Error) -> pyo3::PyErr {
-        pyo3::exceptions::PyValueError::new_err("Error")
+cfg_python! {
+    impl std::convert::From<Error> for pyo3::PyErr {
+        // TODO: Add proper error reporting
+        fn from(_e: Error) -> pyo3::PyErr {
+            pyo3::exceptions::PyValueError::new_err("Error")
+        }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 29907d6..6579c6c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -53,6 +53,9 @@ compile_error!("feature \"wasm\" is only supported for \"wasm32\" targets.");
 #[cfg(all(target_family = "wasm", feature = "python-bindings"))]
 compile_error!("feature \"python-bindings\" is not supported for \"wasm32\" targets.");
 
+#[macro_use]
+pub mod cfg_macros;
+
 pub mod layers;
 
 pub mod errors;
@@ -77,32 +80,32 @@ pub use packet::Packet;
 #[doc(inline)]
 pub use types::{ENCAP_TYPE_ETH, ENCAP_TYPE_LINUX_SLL, ENCAP_TYPE_LINUX_SLL2};
 
-#[cfg(feature = "python-bindings")]
-use pyo3::prelude::*;
+cfg_python! {
+    use pyo3::prelude::*;
 
-/// Python bindings for packet dissection and sculpting in Rust (scalpel)
-#[cfg(feature = "python-bindings")]
-#[pymodule]
-fn scalpel(py: Python, m: &PyModule) -> PyResult<()> {
-    packet::register(py, m)?;
-    Ok(())
+    /// Python bindings for packet dissection and sculpting in Rust (scalpel)
+    #[pymodule]
+    fn scalpel(py: Python, m: &PyModule) -> PyResult<()> {
+        packet::register(py, m)?;
+        Ok(())
+    }
 }
 
-#[cfg(feature = "wasm")]
-use wasm_bindgen::prelude::*;
+cfg_wasm! {
+    use wasm_bindgen::prelude::*;
 
-#[cfg(feature = "wasm")]
-#[wasm_bindgen]
-pub fn dissect_packet(packet: String) -> String {
-    let _ = layers::register_defaults();
+    #[wasm_bindgen]
+    pub fn dissect_packet(packet: String) -> String {
+        let _ = layers::register_defaults();
 
-    let packet = hex::decode(packet);
+        let packet = hex::decode(packet);
 
-    let packet = packet.unwrap();
+        let packet = packet.unwrap();
 
-    let p = Packet::from_bytes(&packet, ENCAP_TYPE_ETH);
+        let p = Packet::from_bytes(&packet, ENCAP_TYPE_ETH);
 
-    let p = p.unwrap();
+        let p = p.unwrap();
 
-    serde_json::to_string_pretty(&p).unwrap()
+        serde_json::to_string_pretty(&p).unwrap()
+    }
 }
diff --git a/src/packet.rs b/src/packet.rs
index 0a8beab..725b75d 100644
--- a/src/packet.rs
+++ b/src/packet.rs
@@ -12,9 +12,6 @@ use crate::errors::Error;
 use crate::types::{EncapType, LayerCreatorFn};
 use crate::Layer;
 
-#[cfg(feature = "python-bindings")]
-use pyo3::prelude::*;
-
 fn get_encap_types_map() -> &'static RwLock<HashMap<EncapType, LayerCreatorFn>> {
     static ENCAP_TYPES_MAP: OnceLock<RwLock<HashMap<EncapType, LayerCreatorFn>>> = OnceLock::new();
     ENCAP_TYPES_MAP.get_or_init(|| RwLock::new(HashMap::new()))
@@ -176,33 +173,36 @@ impl Packet {
 }
 
 // Python Bindings
-#[allow(clippy::borrow_deref_ref)]
-#[cfg(feature = "python-bindings")]
-#[pymethods]
-impl Packet {
-    #[staticmethod]
-    fn from_bytes_py(bytes: &[u8], encap: EncapType) -> PyResult<Self> {
-        let _ = crate::layers::register_defaults();
+cfg_python! {
+    use pyo3::prelude::*;
 
-        Self::from_bytes(bytes, encap).map_err(|e| e.into())
-    }
+    #[allow(clippy::borrow_deref_ref)]
+    #[pymethods]
+    impl Packet {
+        #[staticmethod]
+        fn from_bytes_py(bytes: &[u8], encap: EncapType) -> PyResult<Self> {
+            let _ = crate::layers::register_defaults();
+
+            Self::from_bytes(bytes, encap).map_err(|e| e.into())
+        }
 
-    fn as_json(&self) -> PyResult<String> {
-        Ok(serde_json::to_string_pretty(self).unwrap())
+        fn as_json(&self) -> PyResult<String> {
+            Ok(serde_json::to_string_pretty(self).unwrap())
+        }
     }
-}
 
-#[cfg(feature = "python-bindings")]
-pub(crate) fn register(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_class::<Packet>()?;
-    Ok(())
+    pub(crate) fn register(_py: Python, m: &PyModule) -> PyResult<()> {
+        m.add_class::<Packet>()?;
+        Ok(())
+    }
 }
 
 #[cfg(test)]
 mod tests {
 
-    #[cfg(feature = "wasm")]
-    use wasm_bindgen_test::wasm_bindgen_test;
+    cfg_wasm! {
+        use wasm_bindgen_test::wasm_bindgen_test;
+    }
 
     use super::*;
     use hex;