From 53340567e3857aa824f613a5daa5c7592ee4ccf1 Mon Sep 17 00:00:00 2001
From: Jonathan LEI <me@xjonathan.dev>
Date: Thu, 18 Jul 2024 05:29:03 +0800
Subject: [PATCH] feat: add `abi` command

---
 Cargo.lock             |  1 +
 Cargo.toml             |  1 +
 src/main.rs            |  3 ++
 src/subcommands/abi.rs | 74 ++++++++++++++++++++++++++++++++++++++++++
 src/subcommands/mod.rs |  3 ++
 5 files changed, 82 insertions(+)
 create mode 100644 src/subcommands/abi.rs

diff --git a/Cargo.lock b/Cargo.lock
index f9ba281..eb67b77 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3805,6 +3805,7 @@ dependencies = [
  "rpassword",
  "serde",
  "serde_json",
+ "serde_json_pythonic",
  "serde_with",
  "sha2",
  "shellexpand",
diff --git a/Cargo.toml b/Cargo.toml
index e429b0c..e229d05 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -35,6 +35,7 @@ regex = "1.8.4"
 rpassword = "7.2.0"
 serde = { version = "1.0.164", features = ["derive"] }
 serde_json = { version = "1.0.99", features = ["preserve_order"] }
+serde_json_pythonic = { version = "0.1.2", default-features = false, features = ["alloc", "raw_value"] }
 serde_with = "2.3.3"
 sha2 = "0.10.8"
 shellexpand = "3.1.0"
diff --git a/src/main.rs b/src/main.rs
index f53025c..c06275c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -59,6 +59,8 @@ enum Subcommands {
     Selector(Selector),
     #[clap(about = "Calculate class hash from any contract artifacts (Sierra, casm, legacy)")]
     ClassHash(ClassHash),
+    #[clap(about = "Extract contract ABI from a class artifact (Sierra or legacy)")]
+    Abi(Abi),
     #[clap(about = "Encode string into felt with the Cairo short string representation")]
     ToCairoString(ToCairoString),
     #[clap(about = "Decode string from felt with the Cairo short string representation")]
@@ -179,6 +181,7 @@ async fn run_command(cli: Cli) -> Result<()> {
         (false, Some(command)) => match command {
             Subcommands::Selector(cmd) => cmd.run(),
             Subcommands::ClassHash(cmd) => cmd.run(),
+            Subcommands::Abi(cmd) => cmd.run(),
             Subcommands::ToCairoString(cmd) => cmd.run(),
             Subcommands::ParseCairoString(cmd) => cmd.run(),
             Subcommands::Mont(cmd) => cmd.run(),
diff --git a/src/subcommands/abi.rs b/src/subcommands/abi.rs
new file mode 100644
index 0000000..ae16927
--- /dev/null
+++ b/src/subcommands/abi.rs
@@ -0,0 +1,74 @@
+use std::path::PathBuf;
+
+use anyhow::Result;
+use clap::Parser;
+use serde_json_pythonic::to_string_pythonic;
+use starknet::core::types::contract::{legacy::LegacyContractClass, CompiledClass, SierraClass};
+
+use crate::{path::ExpandedPathbufParser, utils::print_colored_json};
+
+#[derive(Debug, Parser)]
+pub struct Abi {
+    #[clap(
+        long,
+        help = "Present the ABI as a flattened string in a Pythoic style"
+    )]
+    flatten: bool,
+    #[clap(
+        long,
+        help = "When --flatten is used, serialize the ABI in the Pythoic style instead of compact"
+    )]
+    pythonic: bool,
+    #[clap(
+        value_parser = ExpandedPathbufParser,
+        help = "Path to contract artifact file"
+    )]
+    file: PathBuf,
+}
+
+impl Abi {
+    pub fn run(self) -> Result<()> {
+        if self.pythonic && !self.flatten {
+            anyhow::bail!("--pythonic can only be used with --flatten");
+        }
+
+        // Working around a deserialization bug in `starknet-rs`:
+        //   https://github.com/xJonathanLEI/starknet-rs/issues/392
+
+        if let Ok(class) =
+            serde_json::from_reader::<_, SierraClass>(std::fs::File::open(&self.file)?)
+        {
+            let abi = class.abi;
+            if self.flatten {
+                if self.pythonic {
+                    println!("{}", to_string_pythonic(&abi)?);
+                } else {
+                    println!("{}", serde_json::to_string(&abi)?);
+                }
+            } else {
+                print_colored_json(&abi)?;
+            }
+        } else if let Ok(class) =
+            serde_json::from_reader::<_, LegacyContractClass>(std::fs::File::open(&self.file)?)
+        {
+            let abi = class.abi;
+            if self.flatten {
+                if self.pythonic {
+                    println!("{}", to_string_pythonic(&abi)?);
+                } else {
+                    println!("{}", serde_json::to_string(&abi)?);
+                }
+            } else {
+                print_colored_json(&abi)?;
+            }
+        } else if serde_json::from_reader::<_, CompiledClass>(std::fs::File::open(self.file)?)
+            .is_ok()
+        {
+            anyhow::bail!("cannot extract ABI from casm");
+        } else {
+            anyhow::bail!("failed to parse contract artifact");
+        }
+
+        Ok(())
+    }
+}
diff --git a/src/subcommands/mod.rs b/src/subcommands/mod.rs
index e30069e..c7eeb25 100644
--- a/src/subcommands/mod.rs
+++ b/src/subcommands/mod.rs
@@ -96,3 +96,6 @@ pub use block_traces::BlockTraces;
 
 mod transaction_status;
 pub use transaction_status::TransactionStatus;
+
+mod abi;
+pub use abi::Abi;