@@ -18,7 +18,7 @@ def decode_name(name) -> str:
1818 return bytes (name ).decode ("utf-8" ).strip ()
1919
2020
21- async def generate_spot_configs (drift_client : DriftClient ) -> str :
21+ async def generate_spot_configs (drift_client : DriftClient , env : str ) -> str :
2222 spot_markets = sorted (
2323 drift_client .get_spot_market_accounts (), key = lambda market : market .market_index
2424 )
@@ -34,8 +34,8 @@ async def generate_spot_configs(drift_client: DriftClient) -> str:
3434 )
3535 configs .append (config )
3636
37- output = """
38- mainnet_spot_market_configs : list[SpotMarketConfig] = ["""
37+ output = f """
38+ { env } _spot_market_configs : list[SpotMarketConfig] = ["""
3939
4040 for config in configs :
4141 output += f"""
@@ -51,7 +51,7 @@ async def generate_spot_configs(drift_client: DriftClient) -> str:
5151 return output
5252
5353
54- async def generate_perp_configs (drift_client : DriftClient ) -> str :
54+ async def generate_perp_configs (drift_client : DriftClient , env : str ) -> str :
5555 perp_markets = sorted (
5656 drift_client .get_perp_market_accounts (), key = lambda market : market .market_index
5757 )
@@ -68,8 +68,8 @@ async def generate_perp_configs(drift_client: DriftClient) -> str:
6868 configs .append (config )
6969
7070 # Generate Python code
71- output = """
72- mainnet_perp_market_configs : list[PerpMarketConfig] = ["""
71+ output = f """
72+ { env } _perp_market_configs : list[PerpMarketConfig] = ["""
7373
7474 for config in configs :
7575 output += f"""
@@ -89,25 +89,27 @@ async def generate_perp_configs(drift_client: DriftClient) -> str:
8989async def main ():
9090 parser = argparse .ArgumentParser ()
9191 parser .add_argument ("--market-type" , choices = ["perp" , "spot" ], required = True )
92+ parser .add_argument ("--env" , choices = ["mainnet" , "devnet" ], default = "mainnet" )
9293 args = parser .parse_args ()
9394
94- rpc_url = os .getenv ("MAINNET_RPC_ENDPOINT" )
95+ rpc_env_var = f"{ args .env .upper ()} _RPC_ENDPOINT"
96+ rpc_url = os .getenv (rpc_env_var )
9597 if not rpc_url :
96- raise ValueError ("MAINNET_RPC_ENDPOINT is not set" )
98+ raise ValueError (f" { rpc_env_var } is not set" )
9799
98100 drift_client = DriftClient (
99101 AsyncClient (rpc_url ),
100102 Wallet .dummy (),
101- env = "mainnet" ,
103+ env = args . env ,
102104 account_subscription = AccountSubscriptionConfig ("cached" ),
103105 )
104106
105107 await drift_client .subscribe ()
106108
107109 if args .market_type == "perp" :
108- output = await generate_perp_configs (drift_client )
110+ output = await generate_perp_configs (drift_client , args . env )
109111 else :
110- output = await generate_spot_configs (drift_client )
112+ output = await generate_spot_configs (drift_client , args . env )
111113
112114 print (output )
113115 await drift_client .unsubscribe ()
0 commit comments