linearsolver

#!/usr/bin/env python3

import re
import sys
from sympy import symbols, Eq, solve, sympify

PRECISION = 2

def is_expression(arg):
    return bool(re.match(r".*=-?(\d+(\.\d+)?)$", arg))

def get_constant(expr):
    return re.search(r"(?<==).*", expr).group().strip()

def get_equation(expr):
    eq = re.search(r"^.*(?==)", expr).group().strip()
    return re.sub(r'(\d+)([a-zA-Z]+\d*)', r'\1*\2', eq)

def get_symbols(args):
    variables = set()
    for arg in args:
        vrs = re.findall(r"([a-zA-Z]+\d*)", arg)
        variables.update(vrs)
    return list(variables)

def argument_error(args):
    print("Use:", args[0], "<option> <expr 1> ... <expr N>")
    exit(1)

def main():

    fractions = False
        
    if len(sys.argv) < 2:
        argument_error(sys.argv)

    if not is_expression(sys.argv[1]):
        if sys.argv[1] != "-f":
            argument_error(sys.argv)
        fractions = True
        sys.argv.pop(1)

    syms = get_symbols(sys.argv[1:])
    variables = symbols(syms)
    
    equations = []
    
    for expr in sys.argv[1:]:

        if not is_expression(expr):
            argument_error(sys.argv)

        lhs = sympify(get_equation(expr))
        rhs = sympify(get_constant(expr))
        equations.append(Eq(lhs, rhs))

    sol = solve(equations, variables)

    if not sol:
        print("No solution.")

    elif fractions or len(variables) != len(sys.argv) - 1:
        print(", ".join(f"{var} = {val}" for var, val in sol.items()))

    else:
        sdec = {var: round(float(val.evalf()), PRECISION) for var, val in sol.items()}
        print(", ".join(f"{var} = {val}" for var, val in sdec.items()))

if __name__== "__main__":
    main()