Scripting tricks: Python

Photo by Chris Ried on Unsplash

Lets step up our Python game this time.

Just to keep in mind:

These are python “scripting” tricks and not python “software writing” tricks. So we can get by doing a lot of things here which I wouldn’t do if I am developing software in python.

Starter template

Whenever developing a new script, we can start with this template and populate it as we start writing code.

#! /bin/env/python3
import os
import shutil
import time
import sys
import argparse
import pprint
import logging

def parse_args():
    parser = argparse.ArgumentParser(description='This script does something')
    parser.add_argument('-verbose')
    parser.add_argument('-some_arg',    required=True,  type=str, help="some switch to do something")
    args = parser.parse_args()
    
    return args
def main():
    start = time.time()
    args = parse_args()
logging.getLogger().setLevel(level=logging.DEBUG if args.verbose else logging.INFO)
    logging.basicConfig(format='%(levelname)s - %(message)s')
end = time.time()
    duration = end-start
    logging.debug("Script took {:0.0f} min {:0.0f} seconds".format((duration/60), (duration%60)))
if __name__ == '__main__':
    sys.exit(main())

Read a file into a dictionary

I have to do this a lot of times to read some config file.

This example also shows how to loop through all the lines, how to use startswith, pass, split, strip and format if you didn’t know that already.

in_file  = "input.txt" # input file
name_age = {}  # dictionary to store the file
with open(in_file) as f:
    for line in f:
        if line.startswith("#"):  # ignore comments
            pass
        else:
            (key, val) = line.split(":")  # split line at :
            key = key.strip()  # removing whitespaces
            val = val.strip()
            ival = int(val)  # converting age to int
            name_age[key] = ival
            #we could have directly done this too, but I want to keep it simple
            #name_age[key.strip()] = int(val.strip())
for key, val in name_age.items():
    print("name: {0} age: {1}".format(key, val))

Write into a file

out_file = "out.txt"
fout = open(out_file, "w")  # open in write mode
fout.write("Hello\n")
fout.write("World\n")
fout.close()
fout = open(out_file, "a") # open in append mode
fout.write("Oh hello")
fout.close()

prints this in out.txt

Hello
World
Oh hello

Run system commands or other scripts

import os
my_script = "/some/path/something.pl"
my_cmd = "perl "+ my_script
os.system(my_cmd)
file_to_clean = "output.log"
my_cmd = "rm " + file_to_clean
os.system(my_cmd)

Similarly you can execute a bash scrip or another python script.

Capture command output into a variable

After you run a command like above, you might want to see if the output had an error or a particular string.

You can capture the output into a variable like this

import subprocess
my_cmd = "perl my.pl"
cmd_out = subprocess.check_output(my_cmd,stderr=subprocess.STDOUT, shell=True,universal_newlines=True)
print("Output: {0}".format(cmd_out))

This enables capturing both STDOUT and STDERR and makes sure the new lines are preserved.

Parse command line args

This is a big one because I am putting the complete example here. argparse is an elegant way to take care of this.

Full example code:

Here we are taking care of a lot of things, so you need to read it line by line and understand how we are validating user inputs, how we are printing the help menu and how we are exiting from the script when we encounter a bad argument.

Now lets test it.

When user doesn’t pass arguments

>>python my.py
You need to pass both arguments
usage: my.py -h  prints this help menu
my.py -p <PROJECT> -i <INPUT_CFG_FILE>
This script does something
optional arguments:
  -h, --help   show this help message and exit
  -i CFG_FILE  input cfg file
  -p PROJECT   project name e.g., Hydra, WeaponX

When user passes correct arguments

>>python my.py -p Hydra -i input.txt
Inputs: Hydra input.txt

You can also accept key value pairs as input like this

def parse_args():
    parser = argparse.ArgumentParser(description='Custom perforce label creation script')
    parser.add_argument('-d','--description', required=True,  nargs='*', help="description key:value pairs")
    args = parser.parse_args()

    for k in args.description:
        try:
            k1,v1 = k.split(":")
        except:
            logging.error("descriptions fields must be key:value pairs")
            sys.exit(1)

Print into a log and stdout both

fout= open("my.log", "w")
def printer(text)
    print(text)
    fout.write(text + "\n")
user = "Bobby"
printer("hello {0}".format(user))

Now here we could have opened the file inside our function too, but then it will cause a lot of file i/o operations. We can avoid it by opening it outside and then using the global variable.

Time stuff

from datetime import datetime
from datetime import timedelta
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
work_week = datetime.now().strftime("%YWW%V")
print(timestamp)
print(work_week)

Prints

2020-02-07 13:20
2020WW06

This comes handy when you are creating log files with timestamps and you need a specific time format.

Time Taken by your script

import time
start =  time.time()
...
your code
...
end = time.time()
duration = end - start
print("my script took  {:0.0f} min {:0.0f} seconds".format((duration/60), (duration%60)))

Find the previous and coming Monday's date, based on current date

Run your python script in gvim

This is a nifty trick if you use gvim and would like to quickly run your code rather than switching back and forth from the shell.

Add this line to your ~/.gvimrc

map <F5> <Esc>:w<CR>:!clear;python %<CR>

This will map F5 key to run python. It will also save the file before running it, so you don’t need to worry about it.

Loop through all nested dictionary values?