Towards making smart-importer and beancount-import more cooperative.

beancount_import/reconcile.py

@@ -25,6 +25,7 @@
 
 from .matching import FIXME_ACCOUNT, is_unknown_account, CLEARED_KEY
 
+UNCONFIRMED_ACCOUNT_KEY = 'unconfirmed_account'
 display_prediction_explanation = False
 
 classifier_cache_version_number = 1
@@ -580,7 +581,8 @@ def _get_fixme_transactions(self):
             if isinstance(entry, Transaction):
                 if any(
                         is_unknown_account(posting.account)
-                        for posting in entry.postings):
+                        for posting in entry.postings) \
+                   or self._has_unconfirmed_account(entry):
                     output.append(entry)
         return output
 
@@ -722,18 +724,63 @@ def _get_primary_transaction_amount_number(self, transaction: Transaction):
                 return -source_posting.units.number
         return None
 
+    def _get_unknown_account_names(self, transaction: Transaction):
+        return [posting.account for posting in transaction.postings if posting.meta is not None and NEW_ACCOUNT_KEY in posting.meta]
+
+    def _has_unconfirmed_account(self, transaction: Transaction) -> bool:
+        return any((posting.meta is not None and posting.meta.get(UNCONFIRMED_ACCOUNT_KEY, False))
+                for posting in transaction.postings)
+
+    def _strip_unconfirmed_account_tags(self, transaction: Transaction):
+        '''
+            Strips a transaction of meta tags indicating that a posting had a pre-predicted unconfirmed account.
+            Leaves postings with FIXME account unchanged.
+        '''
+        for posting in transaction.postings:
+            if posting.account == FIXME_ACCOUNT:
+                continue
+            if posting.meta is not None and UNCONFIRMED_ACCOUNT_KEY in posting.meta:
+                posting.meta.pop(UNCONFIRMED_ACCOUNT_KEY)
+
+    def _group_predicted_accounts_by_name(self, transaction: Transaction):
+        '''
+            Takes a list of postings with candidate account names,
+            and groups them into groups that should share the same exact account.
+            Expects each predicted posting to have an UNCONFIRMED_ACCOUNT_KEY meta field.
+        '''
+        num_groups = 0
+        group_numbers = []
+        predicted_account_names = []
+        existing_groups = {}  # type: Dict[str, int]
+        new_accounts = []
+        for posting in transaction.postings:
+            if posting.meta is None or not posting.meta.get(UNCONFIRMED_ACCOUNT_KEY, False):
+                continue
+            group_number = existing_groups.setdefault(posting.account,
+                                                        num_groups)
+            predicted_account_names.append(posting.account)
+            if group_number == num_groups:
+                num_groups += 1
+            group_numbers.append(group_number)
+        return predicted_account_names, group_numbers
+
     def _get_unknown_account_predictions(self,
                                          transaction: Transaction) -> List[str]:
-        group_prediction_inputs = self._feature_extractor.extract_unknown_account_group_features(
-            transaction)
-        group_predictions = [
-            self.predict_account(prediction_input)
-            for prediction_input in group_prediction_inputs
-        ]
-        group_numbers = training.get_unknown_account_group_numbers(transaction)
-        return [
-            group_predictions[group_number] for group_number in group_numbers
-        ]
+        if self._has_unconfirmed_account(transaction): 
+            # if any of the postings have an unconfirmed account, then prediction was handled by smart_importer
+            predicted_account_names, _ = _group_predicted_accounts_by_name(transaction)
+            return predicted_account_names
+        else:
+            group_prediction_inputs = self._feature_extractor.extract_unknown_account_group_features(
+                transaction)
+            group_predictions = [
+                self.predict_account(prediction_input)
+                for prediction_input in group_prediction_inputs
+            ]
+            group_numbers = training.get_unknown_account_group_numbers(transaction)
+            return [
+                group_predictions[group_number] for group_number in group_numbers
+            ]
 
     def _make_candidate_with_substitutions(self,
                                            transaction: Transaction,
@@ -754,8 +801,12 @@ def _make_candidate_with_substitutions(self,
             unique_id: account
             for unique_id, account in zip(unique_ids, new_accounts)
         }
-        group_numbers = training.get_unknown_account_group_numbers(transaction)
-        unknown_names = training.get_unknown_account_names(transaction)
+        if self._has_unconfirmed_account(transaction):
+            _, group_numbers = self._group_predicted_accounts_by_name(transaction)
+            unknown_names = self._get_unknown_account_names(transaction)
+        else:
+            group_numbers = training.get_unknown_account_group_numbers(transaction)
+            unknown_names = training.get_unknown_account_names(transaction)
         substitutions = [
             AccountSubstitution(
                 unique_name=unique_id,
@@ -916,6 +967,7 @@ def accept_candidate(self, candidate: Candidate, ignore=False) -> AcceptCandidat
         for entry in new_entries:
             if isinstance(entry, Transaction):
                 self.posting_db.add_transaction(entry)
+                self._strip_unconfirmed_account_tags(entry)
 
         self._extract_training_examples(new_entries)