write tokens directly to spellbook file (#33)

.env.sample

@@ -12,4 +12,6 @@ ORDERBOOK_HOST=
 ORDERBOOK_USER=
 ORDERBOOK_PORT=
 ORDERBOOK_PASSWORD=
-ORDERBOOK_DB=
+ORDERBOOK_DB=
+
+SPELLBOOK_PATH=

src/missing_tokens.py

@@ -1,6 +1,9 @@
+import fileinput
 import os
 
 from enum import Enum
+from typing import Optional
+
 import web3.exceptions
 from dotenv import load_dotenv
 from dune_client.client import DuneClient
@@ -102,7 +105,16 @@ def fetch_missing_tokens(dune: DuneClient, network: Network) -> list[Address]:
  return [Address(row["token"]) for row in v2_missing.get_rows()]
 
 
-def run_missing_tokens(chain: Network) -> None:
+def replace_line(old_line: str, new_line: str, file_loc: str):
+ """Overwrites old_line with new_line in file at file_loc"""
+ for line in fileinput.input(file_loc, inplace=True):
+ if line == old_line:
+ print(new_line.rstrip("\n"))
+ else:
+ print(line.rstrip("\n"))
+
+
+def run_missing_tokens(chain: Network, insert_loc: Optional[str] = None) -> None:
  """Script's main entry point, runs for given network."""
  w3 = Web3(Web3.HTTPProvider(chain.node_url(os.environ["INFURA_KEY"])))
  client = DuneClient(os.environ["DUNE_API_KEY"])
@@ -128,13 +140,29 @@ def run_missing_tokens(chain: Network) -> None:
  for t in missing_tokens
  if t not in ignored
  )
- print(f"Missing Tokens:\n\n{results}\n")
+ if insert_loc:
+ print(f"Writing Tokens to File {insert_loc}")
+ old_line = " ) AS temp_table (contract_address, symbol, decimals)\n"
+ new_line = " " + results + "\n" + old_line
+ replace_line(old_line, new_line, insert_loc)
+ else:
+ print(f"Missing Tokens:\n\n{results}\n")
  else:
  print(f"No missing tokens detected on {chain}. Have a good day!")
 
 
 if __name__ == "__main__":
  load_dotenv()
+ spellbook_path = os.environ.get("SPELLBOOK_PATH")
  for blockchain in list(Network):
- print(f"Execute on {blockchain}")
- run_missing_tokens(chain=blockchain)
+ chain_name = blockchain.as_dune_v2_repr()
+ print(f"Execute on {chain_name}")
+ token_file = f"models/tokens/{chain_name}/tokens_{chain_name}_erc20.sql"
+ spellbook_file = (
+ os.path.join(spellbook_path, token_file) if spellbook_path else None
+ )
+
+ run_missing_tokens(
+ chain=blockchain,
+ insert_loc=spellbook_file,
+ )

tests/test_missing_tokens.py

@@ -0,0 +1,35 @@
+import os
+import unittest
+
+from src.missing_tokens import replace_line
+
+
+class TestAppendToFile(unittest.TestCase):
+ def test_append_to_file(self):
+ test_file = "./out/test_file.txt"
+ try:
+ os.mkdir("./out")
+ open(test_file, "x")
+ except FileExistsError:
+ pass
+
+ f = open(test_file, "a")
+ file_lines = ["First Line\n", "Second Line\n"]
+ f.write("".join(file_lines))
+ f.close()
+
+ with open(test_file, "r", encoding="utf-8") as file:
+ self.assertEqual(file_lines, file.readlines())
+
+ for i, old_line in enumerate(file_lines):
+ new_line = f"New Line {i + 1}\n"
+ replace_line(old_line, new_line, test_file)
+
+ file_lines[i] = new_line
+ with open(test_file, "r", encoding="utf-8") as file:
+ self.assertEqual(file_lines, file.readlines())
+ os.remove(test_file)
+
+
+if __name__ == "__main__":
+ unittest.main()