Alternative method of simplifying squares of numbers

square_simplification.py

'''
Author: Isa (Seniatical)
License: GNU General Public License (GNU GPL)

A method of calculating the square root of any real number using prime factor trees.
'''
from __future__ import annotations

from typing import Any
from functools import reduce
from operator import mul


class Sqrt(object):
    '''
    A mock object for the results returned from calcuating a root of an object,
    only returned if a number is NOT a perfect square
    '''
    coefficent: float
    input: float
    imaginary: bool

    def __init__(self, coefficient: float, input: float, imaginary: bool = False) -> None:
        self.coefficent = coefficient
        self.input = input
        self.imaginary = imaginary

    @property
    def value(self) -> float:
        r = (self.input ** 0.5) * self.coefficent
        if self.imaginary:
            return r * 1j
        return r
    
    def __repr__(self) -> str:
        imaginary = 'i' if self.imaginary else ''
        if self.coefficent != 1:
            return f'{self.coefficent}*Sqrt({self.input}){imaginary}'
        return f'Sqrt({self.input}){imaginary}'
    
    def __eq__(self, other: Sqrt) -> bool:
        return (self.coefficent == other.coefficent) and (self.input == other.input) and (self.imaginary == other.imaginary)


def prime_factors(factorable: int) -> list[int]:
    i = 2
    factors = []

    while (i ** 2) <= factorable:
        if factorable % i:
            i += 1
        else:
            factorable //= i
            factors.append(i)
    
    if factorable > 1:
        factors.append(factorable)
    return factors


def sqrt(i: Any) -> float | Sqrt:
    if i < 0:
        j = sqrt(-i)
        if isinstance(j, Sqrt):
            return Sqrt(j.coefficent, j.input, True)
        return j * 1j
    
    factors = prime_factors(i)

    ## Group factors in pairs
    pairs = []
    unpaired = []

    for factor in factors:
        if factor in unpaired:
            unpaired.remove(factor)
            pairs.append(factor)
        else:
            unpaired.append(factor)
    
    ## If no unpaired values we have a perfect square
    if not unpaired:
        return reduce(mul, pairs)
    
    pairs = pairs or [1, 1]
    unpaired = unpaired or [1, 1]
    return Sqrt(reduce(mul, pairs), reduce(mul, unpaired))


## Testing
if __name__ == '__main__':

    ## Test sqrt(2)
    ## Should yield : Sqrt(2)
    print(sqrt(2), sqrt(2) == Sqrt(1, 2))

    ## Test sqrt(25)
    ## Should yield : 5
    print(sqrt(25), sqrt(25) == 5)

    ## Test sqrt(50)
    ## Should yield : 5*Sqrt(2)
    print(sqrt(50), sqrt(50) == Sqrt(5, 2))

    ## Test sqrt(-2)
    ## Should yield : Sqrt(2)i
    print(sqrt(-2), sqrt(-2) == Sqrt(1, 2, True))

    ## Test sqrt(-25)
    ## Should yield : 5j
    print(sqrt(-25), sqrt(-25) == 5j)

    ## Test sqrt(-50)
    ## Should yield : 5*Sqrt(2)i
    print(sqrt(-50), sqrt(-50) == Sqrt(5, 2, True))

I got some refactoring again! Also added some docstrings (sourcery generated)

class Sqrt(object):
    '''
    A mock object for the results returned from calcuating a root of an object,
    only returned if a number is NOT a perfect square
    '''
    coefficent: float
    input: float
    imaginary: bool

    def __init__(self, coefficient: float, input: float, imaginary: bool = False) -> None:
        self.coefficent = coefficient
        self.input = input
        self.imaginary = imaginary

    @property
    def value(self) -> float:
        """Returns the value of the square root.

        Calculates the square root of the input multiplied by the coefficient.
        Makes the result imaginary if self.imaginary is True.

        This is a property getter.

        Args:
         self: The SquareRoot instance.

        Returns:
         The value of the square root.
        """

        r = (self.input ** 0.5) * self.coefficent
        return r * 1j if self.imaginary else r
    
    def __repr__(self) -> str:
        imaginary = 'i' if self.imaginary else ''
        if self.coefficent != 1:
            return f'{self.coefficent}*Sqrt({self.input}){imaginary}'
        return f'Sqrt({self.input}){imaginary}'
    
    def __eq__(self, other: Sqrt) -> bool:
        return (self.coefficent == other.coefficent) and (self.input == other.input) and (self.imaginary == other.imaginary)


def prime_factors(factorable: int) -> list[int]:
    """Returns the prime factors of a number.

    Iteratively divides the input number by increasing integers, 
    starting from 2. Each time division results in a remainder of 0, 
    the divisor is recorded as a prime factor. The process continues 
    until the square of the divisor exceeds the number.

    Any remaining value greater than 1 is also appended as a prime factor.

    Args:
     factorable: The number to find prime factors for.

    Returns: 
     A list containing the prime factors of the input number.

    """

    i = 2
    factors = []

    while (i ** 2) <= factorable:
        if factorable % i:
            i += 1
        else:
            factorable //= i
            factors.append(i)
    
    if factorable > 1:
        factors.append(factorable)
    return factors


def sqrt(i: Any) -> float | Sqrt:
    """Calculates the square root of a number.

    Finds the prime factors of the input. 
    Groups the factors in pairs.
    If no unpaired factors, the input is a perfect square.
    Otherwise, returns a SquareRoot object representing the square root.
    Makes the result imaginary if the input is negative.

    Args:
     i: The number to calculate the square root of.

    Returns:
     The square root as a float if a perfect square, otherwise a SquareRoot object.
    """
    
    if i < 0:
        j = sqrt(-i)
        return Sqrt(j.coefficent, j.input, True) if isinstance(j, Sqrt) else j * 1j
    factors = prime_factors(i)

    ## Group factors in pairs
    pairs = []
    unpaired = []

    for factor in factors:
        if factor in unpaired:
            unpaired.remove(factor)
            pairs.append(factor)
        else:
            unpaired.append(factor)

    ## If no unpaired values we have a perfect square
    if not unpaired:
        return reduce(mul, pairs)

    pairs = pairs or [1, 1]
    unpaired = unpaired or [1, 1]
    return Sqrt(reduce(mul, pairs), reduce(mul, unpaired))