Simplifying code - perform mathematical operations based on operator

Question

Here is the code for a calculator in Python:

import time
#Returns the sum of num1 and num2
def add(num1, num2):
    return num1 + num2

#Returns the difference of num1 and num2
def subtract(num1, num2):
    return num1 - num2

#Returns the quotient of num1 and num2
def divide(num1, num2):
    return num1 / num2

#Returns the product of num1 and num2
def multiply(num1, num2):
    return num1 * num2

#Returns the exponentiation of num1 and num2
def power(num1, num2):
    return num1 ** num2

import time

def main ():
    operation = input("What do you want to do? (+, -, *, /, ^): ")
    if(operation != "+" and operation != "-" and operation != "*" and operation != "/" and operation != "^"):
        #invalid operation
        print("You must enter a valid operation")
        time.sleep(3)
    else:
        var1 = int(input("Enter num1: ")) #variable one is identified
        var2 = int(input("Enter num2: ")) #variable two is identified
        if(operation == "+"):
            print (add(var1, var2))
        elif(operation == "-"): 
            print (subtract(var1, var2))
        elif(operation == "/"): 
            print (divide(var1, var2))
        elif(operation == "*"):
            print (multiply(var1, var2))
        else:
            print (power(var1, var2))
main()
input("Press enter to exit")
exit()

About 30 minutes ago I found my old Python folder and took a look at all my basic scripts from 8+ months ago. I found my calculator mini-script and thought it would be fun to recreate it in as few lines as possible (I'm just now learning lambda). Here's what I have:

main = lambda operation,var1,var2: var1+var2 if operation=='+' else var1-var2 if operation=='-' else var1*var2 if operation=='*' else var1/var2 if operation=='/' else 'None'
print(main(input('What operation would you like to perform? [+,-,*,/]: '),int(input('Enter num1: ')),int(input('Enter num2: '))))
input('Press enter to exit')

I know this is a personal question based off of my specific situation, but I would appreciate any help making it shorter. Is there a way to make it more Pythonic? Am I using lambda correctly? Is there a way to handle errors in my shortened version? Any help would be appreciated. I'm very new to this. Thanks!

Answer 1

In order to simplify the code, I will suggest to:

1. Create a function to perform operation taking the help of dictionary.

Note: I shared the alternative with lambda function based on requirement mentioned by user. Personally I would use operator, as it is much cleaner.

Using `operator`:

    import operator

    def perform_operation(my_operator):
        return {
            '+': operator.add,
            '-': operator.sub,
            '*': operator.mul,
            '/': operator.truediv,  # "operator.div" in python 2
            '^': operator.pow,
       }.get(my_operator, '^')  # using `^` as defualt value since in your 
                                # "else" block you are calculating the `pow` 

Using `lambda`:

    def perform_operation(my_operator):
        return {
            '+': lambda x, y: x + y,
            '-': lambda x, y: x - y,
            '*': lambda x, y: x * y,
            '/': lambda x, y: x / float(y), 
            '^': lambda x, y: x ** y,
       }.get(my_operator, '^')  # using `^` as defualt value since in your 
                                # "else" block you are calculating the `pow()`

Sample run:

    >>> perform_operation('/')(3, 5)
    0.6

PS: Looking at the defination you would have got the idea why using operator is more pythonic than lambda

2. Update your else block to make a call to it as:

   var1 = int(input("Enter num1: ")) 
   var2 = int(input("Enter num2: ")) 
   perform_operation(operation)(var1, var2)  # Making call to function created above
   # THE END - nothing more in else block
   

3. Simplify your if condition with:

   if operation not in ["+", "-", "*", "/", "^"]:
       # invalid operation