feat: output format

Author: xubaiwang

rkshare-eastmoney/src/basic_org_info.rs

@@ -8,7 +8,7 @@ use arrow_json::ReaderBuilder;
 use bon::Builder;
 use rkshare_utils::{
     FieldsInfo, Raw, Symbol,
-    data::{Data, Fetch, TypeHint, TypedBytes},
+    data::{Data, Fetch, HasTypeHint, TypeHint, TypedBytes},
     mapping,
 };
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
@@ -123,6 +123,12 @@ pub struct Args<Extra = ()> {
     _extra: PhantomData<Extra>,
 }
 
+impl<E> HasTypeHint for Args<E> {
+    fn type_hint(&self) -> Option<TypeHint> {
+        self.raw.as_ref().map(|_| TypeHint::Json)
+    }
+}
+
 use args_builder::State;
 
 #[allow(deprecated)]

rkshare-eastmoney/src/center_gridlist.rs

@@ -139,6 +139,12 @@ pub struct Args<Extra = ()> {
     _extra: PhantomData<Extra>,
 }
 
+impl<E> HasTypeHint for Args<E> {
+    fn type_hint(&self) -> Option<TypeHint> {
+        self.raw.as_ref().map(|_| TypeHint::Json)
+    }
+}
+
 #[derive(Clone, Debug)]
 #[cfg_attr(feature = "cli", derive(clap::Args))]
 pub struct RawArgs {
@@ -156,7 +162,7 @@ use args_builder::State;
 use bon::Builder;
 use rkshare_utils::{
     FieldsInfo, Raw,
-    data::{Data, Fetch},
+    data::{Data, Fetch, HasTypeHint, TypeHint},
 };
 use serde::{Serialize, de::DeserializeOwned};
 

rkshare-eastmoney/src/cli.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use rkshare_utils::data::{Data, Fetch};
+use rkshare_utils::data::{Data, Fetch, HasTypeHint};
 
 /// 东方财富
 #[derive(clap::Subcommand, Debug)]
@@ -17,3 +17,12 @@ impl Fetch for Eastmoney {
         }
     }
 }
+
+impl HasTypeHint for Eastmoney {
+    fn type_hint(&self) -> Option<rkshare_utils::data::TypeHint> {
+        match self {
+            Eastmoney::BasicOrgInfo(args) => args.type_hint(),
+            Eastmoney::CenterGridlist(args) => args.type_hint(),
+        }
+    }
+}

rkshare-sse/src/cli.rs

@@ -1,4 +1,4 @@
-use rkshare_utils::data::Fetch;
+use rkshare_utils::data::{Fetch, HasTypeHint};
 
 #[derive(clap::Subcommand, Debug)]
 /// 上海证券交易所
@@ -32,3 +32,20 @@ impl Fetch for Stock {
         }
     }
 }
+
+impl HasTypeHint for Stock {
+    fn type_hint(&self) -> Option<rkshare_utils::data::TypeHint> {
+        match self {
+            Stock::Summary(args) => args.type_hint(),
+            Stock::DealDaily(args) => args.type_hint(),
+        }
+    }
+}
+
+impl HasTypeHint for Sse {
+    fn type_hint(&self) -> Option<rkshare_utils::data::TypeHint> {
+        match self {
+            Sse::Stock(stock) => stock.type_hint(),
+        }
+    }
+}

rkshare-sse/src/stock/deal_daily.rs

@@ -5,7 +5,7 @@ use arrow::{array::RecordBatch, datatypes::Schema, json::ReaderBuilder};
 use bon::Builder;
 use rkshare_utils::{
     FieldsInfo, Raw,
-    data::{Data, Fetch, TypeHint, TypedBytes},
+    data::{Data, Fetch, HasTypeHint, TypeHint, TypedBytes},
     mapping,
 };
 use serde::{
@@ -123,3 +123,9 @@ where
         })
     }
 }
+
+impl HasTypeHint for Args {
+    fn type_hint(&self) -> Option<TypeHint> {
+        self.raw.as_ref().map(|_| TypeHint::Json)
+    }
+}

rkshare-sse/src/stock/summary.rs

@@ -5,7 +5,7 @@ use arrow::{array::RecordBatch, datatypes::Schema, json::ReaderBuilder};
 use bon::Builder;
 use rkshare_utils::{
     FieldsInfo, Raw,
-    data::{Data, Fetch, TypeHint, TypedBytes},
+    data::{Data, Fetch, HasTypeHint, TypeHint, TypedBytes},
     mapping,
 };
 use serde::{Serialize, de::DeserializeOwned};
@@ -94,3 +94,9 @@ where
         })
     }
 }
+
+impl HasTypeHint for Args {
+    fn type_hint(&self) -> Option<TypeHint> {
+        self.raw.as_ref().map(|_| TypeHint::Json)
+    }
+}

rkshare-tools/src/cli/get.rs

@@ -1,16 +1,33 @@
+use std::{
+    fs::File,
+    path::{Path, PathBuf},
+};
+
+use arrow::json::writer::JsonArray;
 use rkshare::{
     eastmoney::cli::Eastmoney,
     sse::cli::Sse,
     utils::{
-        data::{Data, Fetch},
+        data::{Data, Fetch, HasTypeHint, TypeHint},
         pretty::pretty_print,
     },
     xueqiu::cli::Xueqiu,
 };
 
-/// 请求数据
+/// 从各平台接口获取数据
+#[derive(clap::Args, Debug)]
+pub struct Get {
+    // TODO: allow specify format
+    /// 输出结果到文件
+    #[arg(long, short, global = true)]
+    output: Option<PathBuf>,
+
+    #[command(subcommand)]
+    command: Command,
+}
+
 #[derive(clap::Subcommand, Debug)]
-pub enum Get {
+enum Command {
     #[command(subcommand)]
     Sse(Sse),
     #[command(subcommand)]
@@ -19,20 +36,72 @@ pub enum Get {
     Xueqiu(Xueqiu),
 }
 
+impl HasTypeHint for Get {
+    fn type_hint(&self) -> Option<TypeHint> {
+        match &self.command {
+            Command::Sse(sse) => sse.type_hint(),
+            Command::Eastmoney(eastmoney) => eastmoney.type_hint(),
+            Command::Xueqiu(xueqiu) => xueqiu.type_hint(),
+        }
+    }
+}
+
 impl Get {
     pub fn action(self) -> anyhow::Result<()> {
-        let data = rt()?.block_on(self.fetch())?;
-        pretty_print(data)?;
+        let type_hint = self.type_hint();
+        let Self { output, command } = self;
+        let format = match &output {
+            Some(path) => Some(
+                OutputFormat::infer_from_path(path).ok_or(anyhow::anyhow!("无法从路径推断格式"))?,
+            ),
+            None => None,
+        };
+        // TODO: early validate
+
+        let data = rt()?.block_on(command.fetch())?;
+
+        match (format, type_hint) {
+            (None, None) => {
+                pretty_print(data)?;
+            }
+            // TODO: should split pretty print
+            (None, Some(_)) => {
+                pretty_print(data)?;
+            }
+            (Some(format), None) => {
+                match format {
+                    // TODO: 输出 JSON
+                    OutputFormat::Json => {
+                        let file = File::create(output.unwrap()).unwrap();
+                        let mut writer =
+                            arrow::json::WriterBuilder::new().build::<_, JsonArray>(file);
+                        writer.write(data.as_arrow().unwrap()).unwrap();
+                    }
+                    // TODO: 输出 CSV
+                    OutputFormat::Csv => {
+                        let file = File::create(output.unwrap()).unwrap();
+                        let mut writer = arrow::csv::WriterBuilder::new()
+                            .with_header(true)
+                            .build(file);
+                        writer.write(data.as_arrow().unwrap()).unwrap();
+                    }
+                }
+            }
+            (Some(_output_format), Some(_data_format)) => {
+                // TODO: 验证类型匹配
+                std::fs::write(output.unwrap(), &data.as_raw().unwrap())?;
+            }
+        }
         Ok(())
     }
 }
 
-impl Fetch for Get {
+impl Fetch for Command {
     async fn fetch(self) -> anyhow::Result<Data> {
         match self {
-            Self::Sse(sse) => sse.fetch().await,
-            Self::Eastmoney(eastmoney) => eastmoney.fetch().await,
-            Self::Xueqiu(xueqiu) => xueqiu.fetch().await,
+            Command::Sse(sse) => sse.fetch().await,
+            Command::Eastmoney(eastmoney) => eastmoney.fetch().await,
+            Command::Xueqiu(xueqiu) => xueqiu.fetch().await,
         }
     }
 }
@@ -42,3 +111,23 @@ fn rt() -> tokio::io::Result<tokio::runtime::Runtime> {
         .enable_all()
         .build()
 }
+
+// TODO: support options like CSV->TSV, json indent and formatting
+/// 输出数据格式。
+pub enum OutputFormat {
+    Json,
+    Csv,
+    // TODO: excel
+    // TODO: parquet
+}
+
+impl OutputFormat {
+    /// 从文件路径推断数据格式。
+    pub fn infer_from_path(path: impl AsRef<Path>) -> Option<Self> {
+        match path.as_ref().extension()?.to_str()? {
+            "json" => Some(Self::Json),
+            "csv" => Some(Self::Csv),
+            _ => None,
+        }
+    }
+}

rkshare-tools/src/cli/mod.rs

@@ -16,7 +16,6 @@ pub struct Cli {
 
 #[derive(clap::Subcommand, Debug)]
 enum Commands {
-    #[command(subcommand)]
     Get(Get),
     Serve(Serve),
     #[command(subcommand)]

rkshare-utils/src/data.rs

@@ -3,6 +3,7 @@
 use std::ops::Deref;
 
 use arrow::array::RecordBatch;
+use bytes::Bytes;
 use derive_more::From;
 
 /// 数据。
@@ -14,6 +15,22 @@ pub enum Data {
     Arrow(RecordBatch),
 }
 
+impl Data {
+    pub fn as_raw(&self) -> Option<&Bytes> {
+        match self {
+            Data::Raw(typed_bytes) => Some(typed_bytes),
+            _ => None,
+        }
+    }
+
+    pub fn as_arrow(&self) -> Option<&RecordBatch> {
+        match self {
+            Data::Arrow(record_batch) => Some(record_batch),
+            _ => None,
+        }
+    }
+}
+
 // impl From<RecordBatch> f
 /// 类型注释
 #[derive(Debug, Clone)]
@@ -61,3 +78,11 @@ pub trait Fetch {
     // TODO: 未来应该构建自己的错误类型
     fn fetch(self) -> impl std::future::Future<Output = anyhow::Result<Data>> + Send;
 }
+
+/// 参数类型提示。
+///
+/// 用于命令行检测输出格式是否正确。
+pub trait HasTypeHint {
+    /// Option 表示结果为 `RecordBatch`.
+    fn type_hint(&self) -> Option<TypeHint>;
+}

rkshare-xueqiu/src/cli.rs

@@ -1,5 +1,5 @@
 use clap::Subcommand;
-use rkshare_utils::data::Fetch;
+use rkshare_utils::data::{Fetch, HasTypeHint};
 
 /// 雪球
 #[derive(Subcommand, Debug)]
@@ -16,3 +16,11 @@ impl Fetch for Xueqiu {
         }
     }
 }
+
+impl HasTypeHint for Xueqiu {
+    fn type_hint(&self) -> Option<rkshare_utils::data::TypeHint> {
+        match self {
+            Xueqiu::Detail(args) => args.type_hint(),
+        }
+    }
+}