cargo fmt

jsai28 · jsai28 · commit a6247706eaec · 2025-04-19T19:28:06.000-06:00
diff --git a/src/bin/main.rs b/src/bin/main.rs
@@ -1,12 +1,15 @@
-use std::env;
+use datafusion::prelude::*;
+use dfkit::commands::{cat, convert, count, describe, dfsplit, query, reverse, schema, sort, view};
 use dfkit::utils::{DfKitError, parse_file_list};
-use structopt::StructOpt;
+use std::env;
 use std::path::PathBuf;
-use datafusion::prelude::*;
-use dfkit::commands::{view, query, convert, describe, schema, count, sort, reverse, dfsplit, cat};
+use structopt::StructOpt;
 
 #[derive(StructOpt, Debug)]
-#[structopt(name = "dfkit", about = "A fast SQL-based CLI tool for working with CSV, Parquet, and JSON data files.")]
+#[structopt(
+    name = "dfkit",
+    about = "A fast SQL-based CLI tool for working with CSV, Parquet, and JSON data files."
+)]
 pub struct Cli {
     #[structopt(subcommand)]
     pub command: Commands,
@@ -38,7 +41,7 @@ pub enum Commands {
         #[structopt(parse(from_os_str))]
         filename: PathBuf,
         #[structopt(parse(from_os_str))]
-        output_filename: PathBuf,
+        output: PathBuf,
     },
 
     #[structopt(about = "Show summary statistics for a file")]
@@ -65,7 +68,7 @@ pub enum Commands {
         filename: PathBuf,
         #[structopt(short, long, use_delimiter = true)]
         columns: Vec<String>,
-        #[structopt(short,long)]
+        #[structopt(short, long)]
         descending: bool,
         #[structopt(short = "o", long = "output", parse(from_os_str))]
         output: Option<PathBuf>,
@@ -83,9 +86,9 @@ pub enum Commands {
     Split {
         #[structopt(parse(from_os_str))]
         filename: PathBuf,
-        #[structopt(short,long)]
+        #[structopt(short, long)]
         chunks: usize,
-        #[structopt(parse(from_os_str))]
+        #[structopt(short, long)]
         output: Option<PathBuf>,
     },
 
@@ -95,9 +98,9 @@ pub enum Commands {
         files: Option<String>,
         #[structopt(long, required_unless = "files")]
         dir: Option<PathBuf>,
-        #[structopt(short, long)]
+        #[structopt(short, long, parse(from_os_str))]
         output: PathBuf,
-    }
+    },
 }
 
 #[tokio::main]
@@ -110,11 +113,18 @@ async fn main() -> Result<(), DfKitError> {
         Commands::View { filename, limit } => {
             view(&ctx, &filename, limit).await?;
         }
-        Commands::Query { filename, sql , output} => {
+        Commands::Query {
+            filename,
+            sql,
+            output,
+        } => {
             query(&ctx, &filename, sql, output).await?;
         }
-        Commands::Convert { filename, output_filename } => {
-            convert(&ctx, &filename, &output_filename).await?;
+        Commands::Convert {
+            filename,
+            output,
+        } => {
+            convert(&ctx, &filename, &output).await?;
         }
         Commands::Describe { filename } => {
             describe(&ctx, &filename).await?;
@@ -125,13 +135,22 @@ async fn main() -> Result<(), DfKitError> {
         Commands::Count { filename } => {
             count(&ctx, &filename).await?;
         }
-        Commands::Sort { filename, columns, descending, output } => {
+        Commands::Sort {
+            filename,
+            columns,
+            descending,
+            output,
+        } => {
             sort(&ctx, &filename, &columns, descending, output).await?;
         }
         Commands::Reverse { filename, output } => {
             reverse(&ctx, &filename, output).await?;
         }
-        Commands::Split { filename, chunks, output} => {
+        Commands::Split {
+            filename,
+            chunks,
+            output,
+        } => {
             let out_dir = output.unwrap_or_else(|| env::current_dir().unwrap());
             dfsplit(&ctx, &filename, chunks, &out_dir).await?;
         }
diff --git a/src/commands.rs b/src/commands.rs
@@ -1,13 +1,17 @@
-use std::fs;
-use std::path::{Path, PathBuf};
-use std::sync::Arc;
+use crate::utils::{DfKitError, file_type, register_table, write_output};
 use datafusion::arrow::compute::concat_batches;
 use datafusion::datasource::MemTable;
 use datafusion::logical_expr::col;
 use datafusion::prelude::SessionContext;
-use crate::utils::{file_type, register_table, write_output, DfKitError};
+use std::fs;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
 
-pub async fn view(ctx: &SessionContext, filename: &Path, limit: Option<usize>) -> Result<(), DfKitError> {
+pub async fn view(
+    ctx: &SessionContext,
+    filename: &Path,
+    limit: Option<usize>,
+) -> Result<(), DfKitError> {
     let df = register_table(&ctx, "t", &filename).await?;
     let limit = limit.unwrap_or(10);
 
@@ -20,7 +24,12 @@ pub async fn view(ctx: &SessionContext, filename: &Path, limit: Option<usize>) -
     Ok(())
 }
 
-pub async fn query(ctx: &SessionContext, filename: &Path, sql: Option<String>, output: Option<PathBuf>) -> Result<(), DfKitError> {
+pub async fn query(
+    ctx: &SessionContext,
+    filename: &Path,
+    sql: Option<String>,
+    output: Option<PathBuf>,
+) -> Result<(), DfKitError> {
     let file_type = file_type(&filename)?;
     let _ = register_table(&ctx, "t", &filename).await?;
     let df_sql = ctx.sql(&*sql.unwrap()).await?;
@@ -35,7 +44,11 @@ pub async fn query(ctx: &SessionContext, filename: &Path, sql: Option<String>, o
     Ok(())
 }
 
-pub async fn convert(ctx: &SessionContext, filename: &Path, output_filename: &Path) -> Result<(), DfKitError> {
+pub async fn convert(
+    ctx: &SessionContext,
+    filename: &Path,
+    output_filename: &Path,
+) -> Result<(), DfKitError> {
     let df = register_table(ctx, "t", &filename).await?;
     let output_file_type = file_type(&output_filename)?;
 
@@ -130,9 +143,16 @@ pub async fn reverse(
     Ok(())
 }
 
-pub async fn dfsplit(ctx: &SessionContext, filename: &Path, chunks: usize, output_dir: &Path) -> Result<(), DfKitError> {
+pub async fn dfsplit(
+    ctx: &SessionContext,
+    filename: &Path,
+    chunks: usize,
+    output_dir: &Path,
+) -> Result<(), DfKitError> {
     if chunks == 0 {
-        return Err(DfKitError::CustomError("Chunks must be greater than 0".into()));
+        return Err(DfKitError::CustomError(
+            "Chunks must be greater than 0".into(),
+        ));
     }
     let df = register_table(ctx, "t", filename).await?;
     let total_rows = df.clone().count().await?;
@@ -144,7 +164,9 @@ pub async fn dfsplit(ctx: &SessionContext, filename: &Path, chunks: usize, outpu
     }
 
     if chunks > total_rows {
-        return Err(DfKitError::CustomError("Chunks must be smaller than total rows".into()));
+        return Err(DfKitError::CustomError(
+            "Chunks must be smaller than total rows".into(),
+        ));
     }
 
     fs::create_dir_all(output_dir)?;
@@ -172,7 +194,11 @@ pub async fn dfsplit(ctx: &SessionContext, filename: &Path, chunks: usize, outpu
     Ok(())
 }
 
-pub async fn cat(ctx: &SessionContext, files: Vec<PathBuf>, out_path: &Path) -> Result<(), DfKitError> {
+pub async fn cat(
+    ctx: &SessionContext,
+    files: Vec<PathBuf>,
+    out_path: &Path,
+) -> Result<(), DfKitError> {
     let mut dfs = vec![];
 
     for (i, file) in files.iter().enumerate() {
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,2 +1,2 @@
-pub mod utils;
 pub mod commands;
+pub mod utils;
diff --git a/src/utils.rs b/src/utils.rs
@@ -1,11 +1,11 @@
-use std::path::{Path, PathBuf};
 use datafusion::arrow::error::ArrowError;
 use datafusion::dataframe::DataFrameWriteOptions;
-use datafusion::prelude::*;
 use datafusion::error::DataFusionError;
-use thiserror::Error;
-use tempfile::NamedTempFile;
+use datafusion::prelude::*;
 use reqwest::Client;
+use std::path::{Path, PathBuf};
+use tempfile::NamedTempFile;
+use thiserror::Error;
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum FileFormat {
@@ -47,21 +47,28 @@ pub enum DfKitError {
     Reqwest(#[from] reqwest::Error),
 }
 
-pub fn file_type(
-    file_path: &Path,
-) -> Result<FileFormat,FileParseError> {
-    match Path::new(file_path).extension().and_then(|ext| ext.to_str()) {
-      Some("csv") => Ok(FileFormat::Csv),
-      Some("parquet") => Ok(FileFormat::Parquet),
-      Some("json") => Ok(FileFormat::Json),
-      Some("avro") => Ok(FileFormat::Avro),
-      Some(_) => Err(FileParseError::UnsupportedFileFormat),
-      None => Err(FileParseError::InvalidExtension),
-  }
+pub fn file_type(file_path: &Path) -> Result<FileFormat, FileParseError> {
+    match Path::new(file_path)
+        .extension()
+        .and_then(|ext| ext.to_str())
+    {
+        Some("csv") => Ok(FileFormat::Csv),
+        Some("parquet") => Ok(FileFormat::Parquet),
+        Some("json") => Ok(FileFormat::Json),
+        Some("avro") => Ok(FileFormat::Avro),
+        Some(_) => Err(FileParseError::UnsupportedFileFormat),
+        None => Err(FileParseError::InvalidExtension),
+    }
 }
 
-pub async fn register_table(ctx: &SessionContext, table_name: &str, file_path: &Path) -> Result<DataFrame, DfKitError> {
-    let path_str = file_path.to_str().ok_or(DfKitError::FileParse(FileParseError::InvalidExtension))?;
+pub async fn register_table(
+    ctx: &SessionContext,
+    table_name: &str,
+    file_path: &Path,
+) -> Result<DataFrame, DfKitError> {
+    let path_str = file_path
+        .to_str()
+        .ok_or(DfKitError::FileParse(FileParseError::InvalidExtension))?;
     let is_url = path_str.starts_with("http://") || path_str.starts_with("https://");
 
     let actual_path = if is_url {
@@ -72,20 +79,38 @@ pub async fn register_table(ctx: &SessionContext, table_name: &str, file_path: &
     };
 
     let file_format = file_type(&actual_path)?;
-    let file_name = actual_path.to_str().ok_or(DfKitError::FileParse(FileParseError::InvalidExtension))?;
+    let file_name = actual_path
+        .to_str()
+        .ok_or(DfKitError::FileParse(FileParseError::InvalidExtension))?;
     match file_format {
-        FileFormat::Csv => ctx.register_csv(table_name, file_name, CsvReadOptions::default()).await?,
-        FileFormat::Parquet => ctx.register_parquet(table_name, file_name, ParquetReadOptions::default()).await?,
-        FileFormat::Json => ctx.register_json(table_name, file_name, NdJsonReadOptions::default()).await?,
-        FileFormat::Avro => ctx.register_avro(table_name, file_name, AvroReadOptions::default()).await?,
+        FileFormat::Csv => {
+            ctx.register_csv(table_name, file_name, CsvReadOptions::default())
+                .await?
+        }
+        FileFormat::Parquet => {
+            ctx.register_parquet(table_name, file_name, ParquetReadOptions::default())
+                .await?
+        }
+        FileFormat::Json => {
+            ctx.register_json(table_name, file_name, NdJsonReadOptions::default())
+                .await?
+        }
+        FileFormat::Avro => {
+            ctx.register_avro(table_name, file_name, AvroReadOptions::default())
+                .await?
+        }
     };
 
     Ok(ctx.table(table_name).await?)
 }
 
-pub fn parse_file_list(files: Option<String>, dir: Option<PathBuf>) -> Result<Vec<PathBuf>, DfKitError> {
+pub fn parse_file_list(
+    files: Option<String>,
+    dir: Option<PathBuf>,
+) -> Result<Vec<PathBuf>, DfKitError> {
     if let Some(file_str) = files {
-        Ok(file_str.split(',')
+        Ok(file_str
+            .split(',')
             .map(|s| PathBuf::from(s.trim()))
             .collect())
     } else if let Some(dir_path) = dir {
@@ -102,17 +127,46 @@ pub fn parse_file_list(files: Option<String>, dir: Option<PathBuf>) -> Result<Ve
         }
         Ok(file_list)
     } else {
-        Err(DfKitError::CustomError("No files or directory provided".into()))
+        Err(DfKitError::CustomError(
+            "No files or directory provided".into(),
+        ))
     }
 }
 
-pub async fn write_output(df: DataFrame, out_path: &Path, format: &FileFormat) -> Result<(), DfKitError> {
+pub async fn write_output(
+    df: DataFrame,
+    out_path: &Path,
+    format: &FileFormat,
+) -> Result<(), DfKitError> {
     match format {
-        FileFormat::Csv => df.write_csv(out_path.to_str().unwrap(), DataFrameWriteOptions::default(), None).await?,
-        FileFormat::Parquet => df.write_parquet(out_path.to_str().unwrap(), DataFrameWriteOptions::default(), None).await?,
-        FileFormat::Json => df.write_json(out_path.to_str().unwrap(), DataFrameWriteOptions::default(), None).await?,
+        FileFormat::Csv => {
+            df.write_csv(
+                out_path.to_str().unwrap(),
+                DataFrameWriteOptions::default(),
+                None,
+            )
+            .await?
+        }
+        FileFormat::Parquet => {
+            df.write_parquet(
+                out_path.to_str().unwrap(),
+                DataFrameWriteOptions::default(),
+                None,
+            )
+            .await?
+        }
+        FileFormat::Json => {
+            df.write_json(
+                out_path.to_str().unwrap(),
+                DataFrameWriteOptions::default(),
+                None,
+            )
+            .await?
+        }
         FileFormat::Avro => {
-            return Err(DfKitError::DataFusion(DataFusionError::NotImplemented("Avro write not supported".into())));
+            return Err(DfKitError::DataFusion(DataFusionError::NotImplemented(
+                "Avro write not supported".into(),
+            )));
         }
     };
     Ok(())
@@ -122,13 +176,17 @@ pub async fn download_to_tempfile(url: &str) -> Result<(NamedTempFile, PathBuf),
     let response = Client::new().get(url).send().await?.bytes().await?;
 
     // Try to extract the file extension from the URL
-    let ext = url.split('.').last().and_then(|e| {
-        let e = e.split('?').next().unwrap_or(e); // strip query string
-        match e {
-            "csv" | "json" | "parquet" | "avro" => Some(e),
-            _ => None,
-        }
-    }).ok_or(FileParseError::InvalidExtension)?;
+    let ext = url
+        .split('.')
+        .last()
+        .and_then(|e| {
+            let e = e.split('?').next().unwrap_or(e); // strip query string
+            match e {
+                "csv" | "json" | "parquet" | "avro" => Some(e),
+                _ => None,
+            }
+        })
+        .ok_or(FileParseError::InvalidExtension)?;
 
     // Create temp file with extension
     let tempfile = NamedTempFile::new()?;
@@ -140,4 +198,3 @@ pub async fn download_to_tempfile(url: &str) -> Result<(NamedTempFile, PathBuf),
 
     Ok((tempfile, path_with_ext))
 }
-
diff --git a/tests/test.rs b/tests/test.rs
diff --git a/tests/test_utils.rs b/tests/test_utils.rs

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-pub mod utils;`
`2`	`1`	`pub mod commands;`
	`2`	`+pub mod utils;`