python_snippets/trading.py at master · barleyj/python_snippets · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
#!/usr/bin/env python

import unittest


stock_distribution = .75
price = 5.00


def handle_transaction(price, account):
    global stock_distribution

    stock_value = price * account['stock']
    current_distribution = stock_value / (account['cash'] + stock_value)
    desired = (stock_value + account['cash']) * stock_distribution

    if current_distribution < stock_distribution:
        stock_to_purchase = int(desired / price) - account['stock']
        account = {'cash': account['cash'] - (stock_to_purchase * price),
                   'stock': account['stock'] + stock_to_purchase
                   }
    if current_distribution > stock_distribution:
        stock_to_sell = account['stock'] - int(desired / price)
        account = {'cash': account['cash'] + (stock_to_sell * price),
                   'stock': account['stock'] - stock_to_sell
                   }

    return account


class TestHandleTransaction(unittest.TestCase):
    def test_purchase(self):
        account = {'cash': 1000.00,
                   'stock': 0
                   }

        result = handle_transaction(price, account)

        self.assertTrue(result['cash'], 250.00)
        self.assertTrue(result['stock'], 150)

    def test_sell(self):
        account = {'cash': 0.00,
                   'stock': 200
                   }

        result = handle_transaction(price, account)

        self.assertEqual(result['cash'], 250.00)
        self.assertEqual(result['stock'], 150)

    def test_buy_rebalance(self):
        account = {'cash': 500.00,
                   'stock': 100
                   }

        result = handle_transaction(price, account)

        self.assertEqual(result['cash'], 250.00)
        self.assertEqual(result['stock'], 150)

    def test_sell_rebalance(self):
        account = {'cash': 125.00,
                   'stock': 175
                   }

        result = handle_transaction(price, account)

        self.assertEqual(result['cash'], 250.00)
        self.assertEqual(result['stock'], 150)


if __name__ == '__main__':
    unittest.main()