Add iter method to store proc-macro.

paritytech · Aug 4, 2020 · 54703ec · 54703ec
1 parent 694a396
commit 54703ec
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 11 deletions.
diff --git a/examples/fetch_all_accounts.rs b/examples/fetch_all_accounts.rs
@@ -14,23 +14,20 @@
 // You should have received a copy of the GNU General Public License
 // along with substrate-subxt.  If not, see <http://www.gnu.org/licenses/>.
 
-use sp_keyring::AccountKeyring;
 use substrate_subxt::{
-    system,
-    system::*,
+    system::AccountStoreExt,
     ClientBuilder,
     DefaultNodeRuntime,
-    PairSigner,
 };
 
 #[async_std::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     env_logger::init();
 
     let client = ClientBuilder::<DefaultNodeRuntime>::new().build().await?;
-    let mut iter = client.accounts_iter();
+    let mut iter = client.account_iter(None);
     while let Some((key, account)) = iter.next().await? {
-        println!("{}: {}", key, account);
+        println!("{:?}: {}", key, account.data.free);
     }
     Ok(())
 }
diff --git a/proc-macro/src/store.rs b/proc-macro/src/store.rs
@@ -73,6 +73,7 @@ pub fn store(s: Structure) -> TokenStream {
     let module = utils::module_name(generics);
     let store_name = utils::ident_to_name(ident, "Store").to_camel_case();
     let store = format_ident!("{}", store_name.to_snake_case());
+    let store_iter = format_ident!("{}_iter", store_name.to_snake_case());
     let store_trait = format_ident!("{}StoreExt", store_name);
     let bindings = utils::bindings(&s);
     let fields = utils::fields(&bindings);
@@ -139,13 +140,19 @@ pub fn store(s: Structure) -> TokenStream {
         }
 
         /// Store extension trait.
-        pub trait #store_trait<T: #module> {
+        pub trait #store_trait<T: #subxt::Runtime + #module> {
             /// Retrieve the store element.
             fn #store<'a>(
                 &'a self,
                 #args
                 hash: Option<T::Hash>,
             ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<#ret, #subxt::Error>> + Send + 'a>>;
+
+            /// Iterate over the store element.
+            fn #store_iter<'a>(
+                &'a self,
+                hash: Option<T::Hash>,
+            ) -> #subxt::KeyIter<T, #ident<#(#params),*>>;
         }
 
         impl<T: #subxt::Runtime + #module> #store_trait<T> for #subxt::Client<T> {
@@ -157,6 +164,13 @@ pub fn store(s: Structure) -> TokenStream {
                 let #marker = core::marker::PhantomData::<T>;
                 Box::pin(async move { self.#fetch(&#build_struct, hash).await })
             }
+
+            fn #store_iter<'a>(
+                &'a self,
+                hash: Option<T::Hash>,
+            ) -> #subxt::KeyIter<T, #ident<#(#params),*>> {
+                self.iter(hash)
+            }
         }
     }
 }
@@ -202,13 +216,18 @@ mod tests {
             }
 
             /// Store extension trait.
-            pub trait AccountStoreExt<T: Balances> {
+            pub trait AccountStoreExt<T: substrate_subxt::Runtime + Balances> {
                 /// Retrieve the store element.
                 fn account<'a>(
                     &'a self,
                     account_id: &'a <T as System>::AccountId,
                     hash: Option<T::Hash>,
                 ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<AccountData<T::Balance>, substrate_subxt::Error>> + Send + 'a>>;
+                /// Iterate over the store element.
+                fn account_iter<'a>(
+                    &'a self,
+                    hash: Option<T::Hash>,
+                ) -> substrate_subxt::KeyIter<T, AccountStore<'a, T>>;
             }
 
             impl<T: substrate_subxt::Runtime + Balances> AccountStoreExt<T> for substrate_subxt::Client<T> {
@@ -219,7 +238,14 @@ mod tests {
                 ) -> core::pin::Pin<Box<dyn core::future::Future<Output = Result<AccountData<T::Balance>, substrate_subxt::Error>> + Send + 'a>>
                 {
                     let _ = core::marker::PhantomData::<T>;
-                    Box::pin(self.fetch_or_default(AccountStore { account_id, }, hash))
+                    Box::pin(async move { self.fetch_or_default(&AccountStore { account_id, }, hash).await })
+                }
+
+                fn account_iter<'a>(
+                    &'a self,
+                    hash: Option<T::Hash>,
+                ) -> substrate_subxt::KeyIter<T, AccountStore<'a, T>> {
+                    self.iter(hash)
                 }
             }
         };

diff --git a/proc-macro/tests/balances.rs b/proc-macro/tests/balances.rs
@@ -89,8 +89,8 @@ subxt_test!({
     account: Alice,
     step: {
         state: {
-            alice: AccountStore { account_id: &alice },
-            bob: AccountStore { account_id: &bob },
+            alice: &AccountStore { account_id: &alice },
+            bob: &AccountStore { account_id: &bob },
         },
         call: TransferCall {
             to: &bob,

diff --git a/src/lib.rs b/src/lib.rs
@@ -573,6 +573,7 @@ mod tests {
                 SubxtClient::from_config(config, test_node::service::new_full)
                     .expect("Error creating subxt client"),
             )
+            .set_page_size(2)
             .build()
             .await
             .expect("Error creating client");
@@ -696,4 +697,15 @@ mod tests {
             .unwrap();
         assert_eq!(keys.len(), 4)
     }
+
+    #[async_std::test]
+    async fn test_iter() {
+        let (client, _) = test_client().await;
+        let mut iter = client.iter::<system::AccountStore<_>>(None);
+        let mut i = 0;
+        while let Some(_) = iter.next().await.unwrap() {
+            i += 1;
+        }
+        assert_eq!(i, 4);
+    }
 }