feat: add identification based on filename to schwab importers

redstreet · Feb 9, 2024 · c33e02f · c33e02f
1 parent 78f9182
commit c33e02f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/beancount_reds_importers/importers/schwab/schwab_csv_brokerage.py b/beancount_reds_importers/importers/schwab/schwab_csv_brokerage.py
@@ -1,5 +1,6 @@
 """ Schwab Brokerage .csv importer."""
 
+import re
 from beancount_reds_importers.libreader import csvreader
 from beancount_reds_importers.libtransactionbuilder import investments
 
@@ -51,6 +52,10 @@ def custom_init(self):
             'Cash In Lieu':                 'cash',
             }
 
+    def deep_identify(self, file):
+        last_three = self.config.get('account_number', '')[-3:]
+        return re.match(self.header_identifier, file.head()) and f'XX{last_three}' in file.name
+
     def skip_transaction(self, ot):
         return ot.type in ['', 'Journal']
 

diff --git a/beancount_reds_importers/importers/schwab/schwab_csv_checking.py b/beancount_reds_importers/importers/schwab/schwab_csv_checking.py
@@ -30,6 +30,10 @@ def custom_init(self):
         }
         self.skip_transaction_types = ['Journal']
 
+    def deep_identify(self, file):
+        last_three = self.config.get('account_number', '')[-3:]
+        return self.column_labels_line in file.head() and f'XX{last_three}' in file.name
+
     def prepare_table(self, rdr):
         rdr = rdr.addfield('amount',
                            lambda x: "-" + x['Withdrawal'] if x['Withdrawal'] != '' else x['Deposit'])