Continued fractions and their convergents

from decimal import *
from math import floor, sqrt
getcontext().prec = 200


def continued_fraction_of_sqrt_n(n):
    ret_list = []
    x1 = Decimal(n).sqrt()
    for i in range(10):  # number of values to return
        t1 = floor(x1)
        d = Decimal(x1 - t1)
        ret_list.append(t1)
        if d == 0:
            return ret_list
        else:
            x1 = Decimal(1) / d
    return ret_list


assert continued_fraction_of_sqrt_n(4) == [2]
# This is exact square root, hence it stops before 10 entries
# sqrt(4) = 2
assert continued_fraction_of_sqrt_n(5) == [2, 4, 4, 4, 4, 4, 4, 4, 4, 4]
# This could be infinte list, but I am returning only first 10
# sqrt(5) = <2> +       1
#                 -------------
#                   <4> +    1
#                         -------
#                         <4> +  1
#                               ---
#                                .(so on)
assert continued_fraction_of_sqrt_n(6) == [2, 2, 4, 2, 4, 2, 4, 2, 4, 2]


def convergents(e):
    n = []  # Numerator
    d = []  # Denominator
    ret_list = []  # Decimals after
    for i, val in enumerate(e):
        if i == 0:
            ni = val
            di = 1
        elif i == 1:
            ni = val * e[i - 1] + 1
            di = val
        else:  # i > 1
            ni = val * n[i - 1] + n[i - 2]
            di = val * d[i - 1] + d[i - 2]

        n.append(ni)
        d.append(di)
    for i, j in zip(n, d):
        ret_list.append(Decimal(i) / Decimal(j))
    return ret_list


assert convergents(continued_fraction_of_sqrt_n(4)) == [2]
print(convergents(continued_fraction_of_sqrt_n(5)))
# A list of numbers slowly converging to sqrt(5) = 2.23606797749979... something
# Basically, it finds value of continued fraction upto depth i
# 2,
# 2 + 1/4
# 2 + 1 / ( 4 + 1/4 )
# 2 + 1 / ( 4 + 1 / ( 4 + 1/4 ) )
#
# Here is the output
# [2,
# 2.25,
# 2.2352941176470588235294117647058823529411764705882352941176470588235294117647058823529411764705882352941176470588235294117647058823529411764705882352941176470588235294117647058823529411764705882352941,
# 2.2361111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111,
# 2.2360655737704918032786885245901639344262295081967213114754098360655737704918032786885245901639344262295081967213114754098360655737704918032786885245901639344262295081967213114754098360655737704918033,
# 2.2360681114551083591331269349845201238390092879256965944272445820433436532507739938080495356037151702786377708978328173374613003095975232198142414860681114551083591331269349845201238390092879256965944,
# 2.2360679700347158779462817467568061392289420792983738351909373287045496071624337657591814361410560935501553078750228393933857116754978987758085145258541933126256166636214142152384432669468298921980632,
# 2.2360679779158040027605244996549344375431331953071083505866114561766735679779158040027605244996549344375431331953071083505866114561766735679779158040027605244996549344375431331953071083505866114561767,
# 2.2360679774766060137054648759278681179932592735899968434664847417242818886252787422741296622509138673644981620829048254233318738608477838080013033428708163202965105031107128674561394576871773462717266,
# 2.2360679775010816787654439690399500024037305898754867554444497860679775010816787654439690399500024037305898754867554444497860679775010816787654439690399500024037305898754867554444497860679775010816788,
# ]
# These are better and better estimates of Q/q = sqrt(n)