From 8c35d355e56552c8fda9a1bf0972e065b81474a3 Mon Sep 17 00:00:00 2001
From: Lex Berezhny <lex@damoti.com>
Date: Mon, 9 Jul 2018 09:55:07 -0400
Subject: [PATCH] added is_claim,is_support,is_update and claim_name attributes
 to txo able

---
 lbrynet/wallet/account.py  |  8 ++++++++
 lbrynet/wallet/database.py | 30 +++++++++++++++++++++++++++++-
 lbrynet/wallet/script.py   |  8 ++++----
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/lbrynet/wallet/account.py b/lbrynet/wallet/account.py
index ed6783109..3f1630cfb 100644
--- a/lbrynet/wallet/account.py
+++ b/lbrynet/wallet/account.py
@@ -21,3 +21,11 @@ class Account(BaseAccount):
 
     def get_certificate(self, claim_id):
         return self.certificates[claim_id]
+
+    def get_balance(self, include_claims=False):
+        if include_claims:
+            return super(Account, self).get_balance()
+        else:
+            return super(Account, self).get_balance(
+                is_claim=0, is_update=0, is_support=0
+            )
diff --git a/lbrynet/wallet/database.py b/lbrynet/wallet/database.py
index 8f5e58ac0..c402c0c03 100644
--- a/lbrynet/wallet/database.py
+++ b/lbrynet/wallet/database.py
@@ -3,9 +3,37 @@ from torba.basedatabase import BaseDatabase
 
 class WalletDatabase(BaseDatabase):
 
+    CREATE_TXO_TABLE = """
+        create table if not exists txo (
+            txoid integer primary key,
+            txhash blob references tx,
+            address blob references pubkey_address,
+            position integer not null,
+            amount integer not null,
+            script blob not null,
+            is_reserved boolean not null default 0,
+            
+            claim_name text,
+            is_claim boolean not null default 0,
+            is_update boolean not null default 0,
+            is_support boolean not null default 0
+        );
+    """
+
     CREATE_TABLES_QUERY = (
             BaseDatabase.CREATE_TX_TABLE +
             BaseDatabase.CREATE_PUBKEY_ADDRESS_TABLE +
-            BaseDatabase.CREATE_TXO_TABLE +
+            CREATE_TXO_TABLE +
             BaseDatabase.CREATE_TXI_TABLE
     )
+
+    def txo_to_row(self, tx, address, txo):
+        row = super(WalletDatabase, self).txo_to_row(tx, address, txo)
+        row.update({
+            'is_claim': txo.script.is_claim_name,
+            'is_update': txo.script.is_update_claim,
+            'is_support': txo.script.is_support_claim,
+        })
+        if txo.script.is_claim_involved:
+            row['claim_name'] = txo.script.values['claim_name']
+        return row
diff --git a/lbrynet/wallet/script.py b/lbrynet/wallet/script.py
index 1d8137c4b..e147dc221 100644
--- a/lbrynet/wallet/script.py
+++ b/lbrynet/wallet/script.py
@@ -67,14 +67,14 @@ class OutputScript(BaseOutputScript):
     def is_claim_name(self):
         return self.template.name.startswith('claim_name+')
 
-    @property
-    def is_support_claim(self):
-        return self.template.name.startswith('support_claim+')
-
     @property
     def is_update_claim(self):
         return self.template.name.startswith('update_claim+')
 
+    @property
+    def is_support_claim(self):
+        return self.template.name.startswith('support_claim+')
+
     @property
     def is_claim_involved(self):
         return self.is_claim_name or self.is_support_claim or self.is_update_claim