add empty chunk handling, importing python crate dependencies, tx filtering by address, support for new datatypes in python crate

crates/cli/src/args.rs

@@ -196,6 +196,10 @@ pub struct Args {
     #[arg(long, help_heading = "Output Options")]
     pub no_report: bool,
 
+    /// Skip writing to dataframe if the df is empty
+    #[arg(long, help_heading = "Acquisition Options")]
+    pub skip_empty: Option<bool>,
+
     /// Address(es)
     #[arg(long, help_heading = "Dataset-specific Options", num_args(1..))]
     pub address: Option<Vec<String>>,

crates/cli/src/parse/execution.rs

@@ -17,6 +17,7 @@ pub(crate) fn parse_execution_env(args: &Args, n_tasks: u64) -> Result<Execution
         .verbose(verbose)
         .report(!args.no_report)
         .report_dir(args.report_dir.clone())
+        .skip_empty(args.skip_empty)
         .args(args_str);
 
     let builder = if !args.no_verbose {

crates/freeze/src/datasets/transactions.rs

@@ -55,7 +55,7 @@ impl Dataset for Transactions {
     }
 
     fn optional_parameters() -> Vec<Dim> {
-        vec![Dim::FromAddress, Dim::ToAddress]
+        vec![Dim::Address, Dim::FromAddress, Dim::ToAddress]
     }
 }
 
@@ -65,8 +65,22 @@ pub type TransactionAndReceipt = (Transaction, Option<TransactionReceipt>);
 #[async_trait::async_trait]
 impl CollectByBlock for Transactions {
     type Response = (Block<Transaction>, Vec<TransactionAndReceipt>, bool);
-
     async fn extract(request: Params, source: Arc<Source>, query: Arc<Query>) -> R<Self::Response> {
+        fn get_addresses() -> Vec<H160> {
+            let env = ExecutionEnv::default();
+            let cli_command = env.cli_command.unwrap();
+            if let Some(address_index) = cli_command.iter().position(|arg| arg == "--address") {
+                cli_command[address_index + 1..]
+                    .to_vec()
+                    .iter()
+                    .take_while(|&arg| !arg.starts_with("--"))
+                    .map(|s| s.parse::<H160>().expect("Invalid H160"))
+                    .collect::<Vec<H160>>()
+            } else {
+                Vec::new()
+            }
+        }
+
         let block = source
             .fetcher
             .get_block_with_txs(request.block_number()?)
@@ -78,7 +92,7 @@ impl CollectByBlock for Transactions {
         // filter by from_address
         let from_filter: Box<dyn Fn(&Transaction) -> bool + Send> =
             if let Some(from_address) = &request.from_address {
-                Box::new(move |tx| tx.from.as_bytes() == from_address)
+                Box::new(move |tx| from_address == tx.from.as_bytes())
             } else {
                 Box::new(|_| true)
             };
@@ -89,16 +103,45 @@ impl CollectByBlock for Transactions {
             } else {
                 Box::new(|_| true)
             };
-        let transactions =
-            block.transactions.clone().into_iter().filter(from_filter).filter(to_filter).collect();
+        // filter by addresses (if either the to or from address are in the vector of addresses)
+        let addresses = get_addresses();
+        let addr_filter: Box<dyn Fn(&Transaction) -> bool + Send> =
+            if let Some(address) = &request.address {
+                Box::new(move |tx| {
+                    let to_address_is_in = addresses.contains(&tx.to.unwrap());
+                    let from_address_is_in = addresses.contains(&tx.from);
+                    if addresses.len() == 1 {
+                        to_address_is_in || from_address_is_in
+                    } else {
+                        if !(to_address_is_in && from_address_is_in) {
+                            tx.to.as_ref().map_or(false, |x| x.as_bytes() == address)
+                        } else {
+                            tx.from.as_bytes() == address
+                        }
+                    }
+                })
+            } else {
+                Box::new(|_| true)
+            };
+        let transactions = block
+            .transactions
+            .clone()
+            .into_iter()
+            .filter(from_filter)
+            .filter(to_filter)
+            .filter(addr_filter)
+            .collect();
 
         // 2. collect receipts if necessary
         // if transactions are filtered fetch by set of transaction hashes, else fetch all receipts
         // in block
         let receipts: Vec<Option<_>> =
             if schema.has_column("gas_used") | schema.has_column("success") {
                 // receipts required
-                let receipts = if request.from_address.is_some() || request.to_address.is_some() {
+                let receipts = if request.from_address.is_some() ||
+                    request.to_address.is_some() ||
+                    request.address.is_some()
+                {
                     source.get_tx_receipts(&transactions).await?
                 } else {
                     source.get_tx_receipts_in_block(&block).await?
@@ -108,8 +151,8 @@ impl CollectByBlock for Transactions {
                 vec![None; block.transactions.len()]
             };
 
-        let transactions_with_receips = transactions.into_iter().zip(receipts).collect();
-        Ok((block, transactions_with_receips, query.exclude_failed))
+        let transactions_with_receipts = transactions.into_iter().zip(receipts).collect();
+        Ok((block, transactions_with_receipts, query.exclude_failed))
     }
 
     fn transform(response: Self::Response, columns: &mut Self, query: &Arc<Query>) -> R<()> {

crates/freeze/src/freeze.rs

@@ -177,7 +177,13 @@ async fn freeze_partition(payload: PartitionPayload) -> Result<(), CollectError>
     let dfs = collect_partition(datatype, partition, query, source).await?;
 
     // write dataframes to disk
+    let mut n_rows = 0;
     for (datatype, mut df) in dfs {
+        let df_height = df.height() as u64;
+        if env.skip_empty == Some(true) && df_height == 0 {
+            continue;
+        };
+        n_rows += df_height;
         let path = paths.get(&datatype).ok_or_else(|| {
             CollectError::CollectError("could not get path for datatype".to_string())
         })?;

crates/freeze/src/types/execution.rs

@@ -25,6 +25,8 @@ pub struct ExecutionEnv {
     pub t_end: Option<SystemTime>,
     /// report directory
     pub report_dir: Option<PathBuf>,
+    /// skips writing empty dfs
+    pub skip_empty: Option<bool>,
 }
 
 impl ExecutionEnv {
@@ -67,6 +69,7 @@ pub struct ExecutionEnvBuilder {
     t_start: SystemTime,
     t_end: Option<SystemTime>,
     report_dir: Option<PathBuf>,
+    skip_empty: Option<bool>,
 }
 
 impl Default for ExecutionEnvBuilder {
@@ -82,6 +85,7 @@ impl Default for ExecutionEnvBuilder {
             t_start: SystemTime::now(),
             t_end: None,
             report_dir: None,
+            skip_empty: Some(false),
         }
     }
 }
@@ -134,6 +138,11 @@ impl ExecutionEnvBuilder {
         self
     }
 
+    /// skip_empty
+    pub fn skip_empty(mut self, skip_empty: Option<bool>) -> Self {
+        self.skip_empty = skip_empty;
+        self
+    }
     /// build final output
     pub fn build(self) -> ExecutionEnv {
         ExecutionEnv {
@@ -147,6 +156,7 @@ impl ExecutionEnvBuilder {
             t_start: self.t_start,
             t_end: self.t_end,
             report_dir: self.report_dir,
+            skip_empty: self.skip_empty,
         }
     }
 }

crates/python/pyproject.toml

@@ -5,12 +5,15 @@ build-backend = "maturin"
 [project]
 name = "cryo_python"
 requires-python = ">=3.7"
+dependencies = ["polars","pyarrow"]
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 
+[project.optional-dependencies]
+pandas = ["pandas>=2.1.3"]
 
 [tool.maturin]
 python-source = "python"

crates/python/python/cryo/_args.py

@@ -37,7 +37,8 @@ def parse_cli_args(
     else:
         raise Exception('unknown file_format')
 
-    kwargs['no_verbose'] = not verbose
-
+    if 'no_verbose' not in kwargs.keys():
+        kwargs['no_verbose'] = not verbose
+
     return kwargs
 

crates/python/python/cryo/_spec.py

@@ -54,4 +54,14 @@ class CryoCliArgs(TypedDict, total=False):
         topic3: str | bytes | None
         inner_request_size: int | None
         no_verbose: bool
+        address: typing.Sequence[str] | None
+        to_address: typing.Sequence[str] | None
+        from_address: typing.Sequence[str] | None
+        call_data: typing.Sequence[str] | None
+        function: typing.Sequence[str] | None
+        inputs: typing.Sequence[str] | None
+        slot: typing.Sequence[str] | None
+        event_signature: str | bytes | None
+        inner_request_size: int | None
+        js_tracer: str | bytes | None
 

crates/python/rust/collect_adapter.rs

@@ -64,6 +64,7 @@ use cryo_freeze::collect;
         verbose = false,
         no_verbose = false,
         event_signature = None,
+        skip_empty = Some(false),
     )
 )]
 #[allow(clippy::too_many_arguments)]
@@ -125,6 +126,7 @@ pub fn _collect(
     verbose: bool,
     no_verbose: bool,
     event_signature: Option<String>,
+    skip_empty: Option<bool>,
 ) -> PyResult<&PyAny> {
     if let Some(command) = command {
         pyo3_asyncio::tokio::future_into_py(py, async move {
@@ -190,6 +192,7 @@ pub fn _collect(
             verbose,
             no_verbose,
             event_signature,
+            skip_empty,
         };
         pyo3_asyncio::tokio::future_into_py(py, async move {
             match run_collect(args).await {

crates/python/rust/freeze_adapter.rs

@@ -61,6 +61,7 @@ use cryo_cli::{run, Args};
         verbose = false,
         no_verbose = false,
         event_signature = None,
+        skip_empty = Some(false),
     )
 )]
 #[allow(clippy::too_many_arguments)]
@@ -122,6 +123,7 @@ pub fn _freeze(
     verbose: bool,
     no_verbose: bool,
     event_signature: Option<String>,
+    skip_empty: Option<bool>,
 ) -> PyResult<&PyAny> {
     if let Some(command) = command {
         freeze_command(py, command)
@@ -182,6 +184,7 @@ pub fn _freeze(
             verbose,
             no_verbose,
             event_signature,
+            skip_empty,
         };
 
         pyo3_asyncio::tokio::future_into_py(py, async move {