LLD || Stock Exchange

Please review my code below and give me suggestions to improve.
Requirement:

  1. User will be able to buy/sell stocks
  2. User can make transactions
  3. User can check list of stock they bought
  4. If price of stock goes up/down then it should also reflect in the users account

Models

import random


class Payment:
    def __init__(self, amount, user_id):
        self.txn_id = str(random.random())
        self.amount = amount
        self.user_id = user_id

    def set_amount(self, amount):
        self.amount = amount

    def get_amount(self):
        return self.amount

    def get_user(self):
        return self.user_id

import uuid
from abc import ABC


class Stock(ABC):
    uuid: uuid.UUID
    current_price: int
    name: str
    listing_date: str

    def __init__(self, current_price, name):
        self.uuid = uuid.uuid4()
        self.current_price = current_price
        self.name = name

    def set_price(self, price):
        self.current_price = price

    def get_price(self):
        return self.current_price

    def get_name(self):
        return self.name

    def set_name(self, name):
        self.name = name

    def update_price(self, change):
        self.current_price += change

from collections import defaultdict


from collections import defaultdict


class StockExchange:
    _instance = None

    class OnlyOne:
        def __init__(self, name):
            pass

    def __init__(self, name):
        if not StockExchange._instance:
            self.name = name
            self.stocks = []  # dictionary of stocks
            StockExchange._instance = self.OnlyOne(name)

    def get_listed_companies(self):
        return self.stocks

    def add_stock(self, stock):
        if stock not in self.stocks:
            self.stocks.append(stock)
            return True
        return False
    
    def remove_stock(self, stock):
        if stock in self.stocks:
            self.stocks.pop(stock.uuid)
            return True
        return False


from abc import ABC
from enum import Enum
from collections import defaultdict

import uuid as uuid


class UserType(Enum):
    DEFAULT, ADMIN = 0, 1


class Address:
    city: str
    zipcode: str
    country: str
    state: str


class User:
    id: uuid.UUID
    name: str
    address: Address
    user_type: UserType.DEFAULT

    def __init__(self, name):
        self.id = uuid.uuid4()
        self.name = name
        self.amount = 0
        self.stocks = defaultdict()

    def get_amount(self):
        return self.amount

    def set_amount(self, amount):
        self.amount = amount

    def update_amount(self, amount):
        self.amount += amount

    def get_user_name(self):
        return self.name

    def get_id(self):
        return self.id

    def add_stock(self, stock, qty):
        self.stocks[stock.uuid] = self.stocks.get(stock.uuid, 0) + qty
        print(f'You have {self.stocks[stock.uuid]} total stock of the company {stock.name}')
        return True

    def remove_stock(self, stock, qty):
        if self.stocks.get(stock.uuid, 0) - qty >= 0:
            self.stocks[stock.uuid] -= (self.stocks.get(stock.uuid, 0) + qty)
        else:
            return False
        if self.stocks.get(stock.uuid, 0) == 0:
            self.stocks.pop(stock.uuid)
        total = stock.get_price() * qty
        self.update_amount(total)
        return True

Services

from LLD.StockExchange.models.Payment import Payment
from collections import defaultdict


class PaymentService:
    def __init__(self):
        self.txns = defaultdict(list)

    def make_payment(self, amount, user):
        if user.get_amount() - amount >= 0:
            user.update_amount(-amount)
            self.txns[user.id].append(Payment(amount, user.id))
            return True
        return False

    def get_transactions_history(self, user):
        return self.txns[user.id]

from LLD.StockExchange.services.PaymentService import PaymentService


class TradingService:
    def __init__(self):
        self.current_stock = None
        self.current_user = None
        self.payment_service = PaymentService()

    def buy_stock(self, qty):
        if self.current_stock and self.current_user:
            total_amount = self.current_stock.get_price() * qty  # TODO add amount in the user profile as well and check if its exceeding total amount
            if self.payment_service.make_payment(total_amount, self.current_user):
                self.current_user.add_stock(self.current_stock, qty)
                print(f'Successfully bought stock {self.current_stock.name}')
                return True
            print('You don\'t have sufficient fund!')
            return False
        else:
            print(f'Something went wrong. Invalid current stock or user')
            return False

    def sell_stock(self, qty):
        if self.current_stock and self.current_user:
            return self.current_user.remove_stock(self.current_stock, qty)
        else:
            print(f'Something went wrong. Invalid current stock or user')
            return False

class UserManager:
    def __init__(self):
        self.users = {}

    def add_user(self, user):
        if user.id not in self.users:
            self.users[user.id] = user
            return True
        return False

    def get_user(self, user):
        return self.users[user.id]

    def remove_user(self, user):
        if user.id in self.users:
            self.users.pop(user.id)
            return True
        return False

Testing Our Application

from LLD.StockExchange.models.Stock import Stock
from LLD.StockExchange.models.StockExchange import StockExchange
from LLD.StockExchange.services.TradingService import TradingService
from LLD.StockExchange.services.UserManagementService import UserManager

from LLD.StockExchange.models.User import User

if __name__ == "__main__":
    stock_exchange = StockExchange('NSE')
    stock1, stock2 = Stock(30, 'XYZ'), Stock(50, 'ABC')
    stock_exchange.stocks.append(stock1)
    stock_exchange.stocks.append(stock2)
    user1 = User('ajay')
    um = UserManager()
    um.add_user(user1)
    ts1 = TradingService()
    ts1.current_user = user1
    ts1.current_stock = stock1

    ts1.buy_stock(4)
    ts1.buy_stock(8)
    ts1.sell_stock(2)

    user1.set_amount(500)
    ts1.buy_stock(4)
    ts1.current_stock = stock2
    ts1.buy_stock(8)
    ts1.sell_stock(2)
    print(user1.get_amount())
    print(user1.stocks)

Output

You don't have sufficient fund!
You don't have sufficient fund!
You have 4 total stock of the company XYZ
Successfully bought stock XYZ
You don't have sufficient fund!
380
defaultdict(None, {UUID('da3c3252-0bfa-46c8-96b8-6f9999a93458'): 4})

Comments (2)