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()