Check-after now accounts for funding options.

Check-before will still output the current and expected balance of owners without an associated token account as 'No token account'. Check-after will output the actual balance, regardless of the expected balance it reads from the 'before.csv' file. This is not a perfect solution, but it makes some sense so I will keep it for now. E.g. before: 2te...xf3,No token account, No token account. After: 2te...xf3,No token account, 10.0, NaN
nukec · Jun 8, 2021 · dc930e5 · dc930e5
1 parent b4b43bf
commit dc930e5
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 54 deletions.
diff --git a/tools/flat-distributor/flat-distributor.py b/tools/flat-distributor/flat-distributor.py
@@ -335,8 +335,6 @@ def before(input_file, drop, addr_type, mint, decimals, rpc_url):
     with open(output_file, 'w') as fw:
         for line in lines:
             try:
-                #addr = line.split(',')[0].strip()
-                #balance = float(line.split(',')[1].strip())
                 addr = line.strip()
                 ok, balance = get_balance(addr, addr_type, mint, rpc_url)
             except (IndexError, ValueError) as e:
@@ -357,37 +355,43 @@ def after(input_file, addr_type, mint, decimals, url):
     with open(input_file, 'r') as f:
         lines = f.readlines()
 
-    with open(output_file, 'w') as fw:
+    with open(output_file, 'w') as f:
+        # Read before.csv
         for line in lines:
+            output_line = ''
             try:
-                addr, _, expected = line.split(',')
-                try:
-                    expected = float(expected)
-                except ValueError:
-                    state = f'{addr},{expected},{expected},NaN'
-                    print(state)
-                    fw.write(state + '\n')
-                    continue
-            except (IndexError) as e:
-                sys.exit('Error when reading input file: ' + str(e))
+                addr, _, expected = [x.strip() for x in line.split(',')]
+                output_line += f'{addr},'
+            except IndexError as e:
+                sys.exit('Error reading input file: ' + str(e))
+
+            try:
+                expected = float(expected)
+                output_line += f'{expected:.{decimals}f},'
+            except ValueError:
+                # Not a number, expecting a No token account message
+                output_line += f'{expected},'
 
             ok, actual = get_balance(addr, addr_type, mint, url)
             if ok:
-                diff = actual - expected
                 endc = '\033[0m'
-                if diff >= 0:
-                    # green
-                    color = '\033[92m'
-                elif diff < 0:
-                    # red
-                    color = '\033[91m'
-                state = f'{addr},{expected:.{decimals}f},{actual:.{decimals}f},{diff:f}'
-                print(color + state + endc)
-                fw.write(state + '\n')
+                startc = ''
+                try:
+                    diff = actual - expected
+                    if diff >= 0:
+                        startc = '\033[92m'
+                    else:
+                        startc = '\033[91m'
+                    output_line += f'{actual:.{decimals}f},{diff:.{decimals}f}'
+                except TypeError:
+                    # Assuming actual was not a number
+                    diff = 'NaN'
+                    output_line += f'{actual},{diff}'          
             else:
-                state = f'{addr},{expected},{actual},NaN'
-                print(state)
-                fw.write(state + '\n')
+                output_line += f'{actual},NaN'
+
+            print(startc + output_line + endc)
+            f.write(output_line + '\n')
 
 
 def transfer(input_path, interactive, drop_amount, 

diff --git a/tools/proportional-distributor/proportional-distributor.py b/tools/proportional-distributor/proportional-distributor.py
@@ -355,37 +355,43 @@ def after(input_file, addr_type, mint, decimals, url):
     with open(input_file, 'r') as f:
         lines = f.readlines()
 
-    with open(output_file, 'w') as fw:
-        for line in lines:
-            try:
-                addr, _, expected = line.split(',')
+    with open(output_file, 'w') as f:
+            # Read before.csv
+            for line in lines:
+                output_line = ''
+                try:
+                    addr, _, expected = [x.strip() for x in line.split(',')]
+                    output_line += f'{addr},'
+                except IndexError as e:
+                    sys.exit('Error reading input file: ' + str(e))
+
                 try:
                     expected = float(expected)
+                    output_line += f'{expected:.{decimals}f},'
                 except ValueError:
-                    state = f'{addr},{expected},{expected},NaN'
-                    print(state)
-                    fw.write(state + '\n')
-                    continue
-            except (IndexError) as e:
-                sys.exit('Error when reading input file: ' + str(e))
-
-            ok, actual = get_balance(addr, addr_type, mint, url)
-            if ok:
-                diff = actual - expected
-                if diff >= 0:
-                    # green
-                    color = '\033[92m'
-                elif diff < 0:
-                    # red
-                    color = '\033[91m'
-                state = f'{addr},{expected:.{decimals}f},{actual:.{decimals}f},{diff:f}'
-                endc = '\033[0m'
-                print(color + state + endc)
-                fw.write(state + '\n')
-            else:
-                state = f'{addr},{expected},{actual},NaN'
-                print(state)
-                fw.write(state + '\n')
+                    # Not a number, expecting a No token account message
+                    output_line += f'{expected},'
+
+                ok, actual = get_balance(addr, addr_type, mint, url)
+                if ok:
+                    endc = '\033[0m'
+                    startc = ''
+                    try:
+                        diff = actual - expected
+                        if diff >= 0:
+                            startc = '\033[92m'
+                        else:
+                            startc = '\033[91m'
+                        output_line += f'{actual:.{decimals}f},{diff:.{decimals}f}'
+                    except TypeError:
+                        # Assuming actual was not a number
+                        diff = 'NaN'
+                        output_line += f'{actual},{diff}'          
+                else:
+                    output_line += f'{actual},NaN'
+
+                print(startc + output_line + endc)
+                f.write(output_line + '\n')
 
 
 def transfer(input_path, interactive, drop_amount,