add wallet balance

lnbits/core/crud.py

@@ -31,6 +31,7 @@ from .models import (
     TinyURL,
     User,
     Wallet,
+    WalletBalance,
     WebPushSubscription,
 )
 
@@ -419,7 +420,7 @@ async def delete_unused_wallets(
 
 async def get_wallet(
     wallet_id: str, deleted: Optional[bool] = None, conn: Optional[Connection] = None
-) -> Optional[Wallet]:
+) -> Optional[WalletBalance]:
     where = "AND deleted = :deleted" if deleted is not None else ""
     return await (conn or db).fetchone(
         f"""
@@ -429,13 +430,13 @@ async def get_wallet(
         WHERE id = :wallet {where}
         """,
         {"wallet": wallet_id, "deleted": deleted},
-        Wallet,
+        WalletBalance,
     )
 
 
 async def get_wallets(
     user_id: str, deleted: Optional[bool] = None, conn: Optional[Connection] = None
-) -> list[Wallet]:
+) -> list[WalletBalance]:
     where = "AND deleted = :deleted" if deleted is not None else ""
     return await (conn or db).fetchall(
         f"""
@@ -445,14 +446,14 @@ async def get_wallets(
         WHERE "user" = :user {where}
         """,
         {"user": user_id, "deleted": deleted},
-        Wallet,
+        WalletBalance,
     )
 
 
 async def get_wallet_for_key(
     key: str,
     conn: Optional[Connection] = None,
-) -> Optional[Wallet]:
+) -> Optional[WalletBalance]:
     return await (conn or db).fetchone(
         """
         SELECT *, COALESCE((
@@ -462,7 +463,7 @@ async def get_wallet_for_key(
         WHERE (adminkey = :key OR inkey = :key) AND deleted = false
         """,
         {"key": key},
-        Wallet,
+        WalletBalance,
     )
 
 

lnbits/core/models.py

@@ -42,25 +42,10 @@ class Wallet(BaseModel):
     adminkey: str
     inkey: str
     deleted: bool = False
-    balance_msat: int = 0
     created_at: datetime = datetime.now(timezone.utc)
     updated_at: datetime = datetime.now(timezone.utc)
     currency: Optional[str] = None
 
-    # @property
-    # def balance_msat(self) -> int:
-    #     return self.balance_msat // 1000
-
-    @property
-    def balance(self) -> int:
-        return self.balance_msat // 1000
-
-    @property
-    def withdrawable_balance(self) -> int:
-        from .services import fee_reserve
-
-        return self.balance_msat - fee_reserve(self.balance_msat)
-
     @property
     def lnurlwithdraw_full(self) -> str:
         url = url_for("/withdraw", external=True, usr=self.user, wal=self.id)
@@ -78,6 +63,22 @@ class Wallet(BaseModel):
         )
 
 
+class WalletBalance(Wallet):
+    """Wallet with balance properties"""
+
+    balance_msat: int = 0
+
+    @property
+    def balance(self) -> int:
+        return self.balance_msat // 1000
+
+    @property
+    def withdrawable_balance(self) -> int:
+        from .services import fee_reserve
+
+        return self.balance_msat - fee_reserve(self.balance_msat)
+
+
 class KeyType(Enum):
     admin = 0
     invoice = 1
@@ -91,7 +92,7 @@ class KeyType(Enum):
 @dataclass
 class WalletTypeInfo:
     key_type: KeyType
-    wallet: Wallet
+    wallet: WalletBalance
 
 
 class UserExtra(BaseModel):
@@ -173,7 +174,7 @@ class User(BaseModel):
     username: Optional[str] = None
     pubkey: Optional[str] = None
     extensions: list[str] = []
-    wallets: list[Wallet] = []
+    wallets: list[WalletBalance] = []
     admin: bool = False
     super_user: bool = False
     has_password: bool = False
@@ -275,13 +276,6 @@ class CreatePayment(BaseModel):
     fee: int = 0
 
 
-# TODO: thanks about paymeny extra
-# class PaymentExtra(Extra):
-#     tag: Optional[str] = None
-#     def __getitem__(self, key):
-#         return self[key] or self._raw_json[key]
-
-
 class Payment(BaseModel):
     status: str
     checking_id: str

lnbits/core/services.py

@@ -78,7 +78,7 @@ from .models import (
     PaymentState,
     User,
     UserExtra,
-    Wallet,
+    WalletBalance,
 )
 
 
@@ -452,7 +452,7 @@ async def _create_external_payment(
 
 
 def _check_wallet_balance(
-    wallet: Wallet,
+    wallet: WalletBalance,
     fee_reserve_total_msat: int,
     internal_checking_id: Optional[str] = None,
 ):
@@ -701,7 +701,7 @@ def fee_reserve_total(amount_msat: int, internal: bool = False) -> int:
     return fee_reserve(amount_msat, internal) + service_fee(amount_msat, internal)
 
 
-async def send_payment_notification(wallet: Wallet, payment: Payment):
+async def send_payment_notification(wallet: WalletBalance, payment: Payment):
     await websocket_updater(
         wallet.inkey,
         json.dumps(

lnbits/core/views/generic.py

@@ -13,7 +13,7 @@ from pydantic.types import UUID4
 
 from lnbits.core.extensions.models import Extension, ExtensionMeta, InstallableExtension
 from lnbits.core.helpers import to_valid_user_id
-from lnbits.core.models import User
+from lnbits.core.models import User, WalletBalance
 from lnbits.core.services import create_invoice, create_user_account
 from lnbits.decorators import check_admin, check_user_exists
 from lnbits.helpers import template_renderer
@@ -167,7 +167,8 @@ async def wallet(
     if wal:
         wallet = await get_wallet(wal.hex)
     elif len(user.wallets) == 0:
-        wallet = await create_wallet(user_id=user.id)
+        _wallet = await create_wallet(user_id=user.id)
+        wallet = WalletBalance(**_wallet.dict())
         user.wallets.append(wallet)
     elif lnbits_last_active_wallet and user.get_wallet(lnbits_last_active_wallet):
         wallet = await get_wallet(lnbits_last_active_wallet)

lnbits/core/views/user_api.py

@@ -21,7 +21,7 @@ from lnbits.core.models import (
     AccountOverview,
     CreateTopup,
     User,
-    Wallet,
+    WalletBalance,
 )
 from lnbits.core.services import update_wallet_balance
 from lnbits.db import Filters, Page
@@ -103,7 +103,7 @@ async def api_users_toggle_admin(user_id: str) -> None:
 
 
 @users_router.get("/user/{user_id}/wallet")
-async def api_users_get_user_wallet(user_id: str) -> List[Wallet]:
+async def api_users_get_user_wallet(user_id: str) -> List[WalletBalance]:
     return await get_wallets(user_id)
 
 

lnbits/core/views/wallet_api.py

@@ -12,6 +12,7 @@ from lnbits.core.models import (
     CreateWallet,
     KeyType,
     Wallet,
+    WalletBalance,
 )
 from lnbits.decorators import (
     WalletTypeInfo,
@@ -61,7 +62,7 @@ async def api_update_wallet(
     name: Optional[str] = Body(None),
     currency: Optional[str] = Body(None),
     key_info: WalletTypeInfo = Depends(require_admin_key),
-) -> Wallet:
+) -> WalletBalance:
     wallet = await get_wallet(key_info.wallet.id)
     if not wallet:
         raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Wallet not found")
@@ -81,7 +82,7 @@ async def api_delete_wallet(
     )
 
 
-@wallet_router.post("", response_model=Wallet)
+@wallet_router.post("")
 async def api_create_wallet(
     data: CreateWallet,
     wallet: WalletTypeInfo = Depends(require_admin_key),