Please review my code below and give me suggestions to improve.
Requirement:
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
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})