Learn Python: Automate your routine with a python script

In two of the previous posts, show Git branch on your Linux prompt and Git colored output, we’ve learned how to make work on your Debian/Ubuntu Linux with Git more usable.  It would be useful to have some script that does this for you to save you or someone else some time in the future. My idea is to script it with Python!

This article is a tutorial with an example of the usage of various basic Python modules. You can use parts of it as building blocks for your own automation needs.

Step 1: Create dist files to append .bashrc and .gitconfig

First, we need to create two empty files and put in a few lines of code:

# Add git branch if its present to PS1
parse_git_branch() {
 git branch 2> /dev/null | sed -e '/^[^*]/d' -e 's/* \(.*\)/(\1)/'
}
if [ "$color_prompt" = yes ]; then
 PS1='${debian_chroot:+($debian_chroot)}\[\033[01;32m\]\u@\h\[\033[00m\]:\[\033[01;34m\]\w\[\033[01;31m\] $(parse_git_branch)\[\033[00m\]\$ '
else
 PS1='${debian_chroot:+($debian_chroot)}\u@\h:\w $(parse_git_branch)\$ '
fi

bashrc.dist – to display Git branch name on your command prompt.

[alias]
lg=log --graph --pretty=format:'%Cred%h%Creset -%C(yellow)%d%Creset %s %Cgreen(%cr) %C(bold blue)<%an>%Creset' --abbrev-commit --date=relative
st=status
ci=commit
co=checkout

[color]
branch=auto
diff=auto
status=auto

[color "branch"]
current=yellow reverse
local=yellow
remote=green

[color "diff"]
meta=yellow bold
frag=magenta bold
old=red bold
new=green bold

[color "status"]
added=yellow
changed=green
untracked=cyan

gitconfig.dist – to use shorter commands when operating Git.

Step 2: Display Git branch on command prompt

Let’s write a simple Python script. We will read/write files, search with regular expressions  and install package (colordiff) using three modules: os (miscellaneous operating system interfaces)re (regular expressions) and subprocess (subprocess management) .

The example below displays the first task for script. It opens your local .bashrc (if it exists), executes method show_branch_in_shell , which modifies bashrc content with code we stored in bashrc.dist beforehand. If any content changes, the script will overwrite the existing file.

#!/usr/bin/python -tt

import re
import os
import subprocess

# get user home directory
home_dir = os.path.expanduser('~')
# path to distfiles
bashrc_dist_filename = 'bashrc.dist'
gitconfig_dist_filename = 'gitconfig.dist'
# We will identify if our addon exists by this comment
git_prompt_regexp = r'# Add git branch if its present to PS1'
# flag for content change
content_changed = False

def show_branch_in_shell(bashrc_content):
    global content_changed

    # modify bashrc_conent only if regexp is not match
    if not re.search(git_prompt_regexp, bashrc_content):
        bashrc_dist_handler = open(bashrc_dist_filename, 'r')
        dist_content = bashrc_dist_handler.read()
        bashrc_dist_handler.close()
        bashrc_content = bashrc_content + '\n' + dist_content
        content_changed = True

    return bashrc_content

def main():
    path_bash = os.path.join(home_dir, '.bashrc')
    if not os.path.isfile(path_bash):
        print('bashrc is not found')
        exit(1)

    # read local .bashrc file
    file_handler = open(path_bash, 'rU')
    bashrc_content = file_handler.read()
    file_handler.close()

    # execute changes
    bashrc_content = show_branch_in_shell(basrc_content)

    if content_changed:
        # save changed bashrc
        file_handler = open(path_bash, 'w')
        file_handler.write(bashrc_content)
        file_handler.close()
if __name__ == '__main__':
    main()

Step 3: Enable Git color output

Now we have script with one method implemented. Our next step is to enable color prompt in bashrc. You can do it manually by setting the force_color_prompt parameter value to ‘yes’. Let’s modify our script. It is good to use regular expressions to find if it is enabled or not:

def color_prompt_enable(bashrc_content):
    # Searches for commented out "force_color_prompt" and enables it
    return re.sub(r'#\s*force_color_prompt\s*=\s*yes', r'force_color_prompt=yes', bashrc_content)

After that, we need to have colordiff package installed on our Linux. Here is implementation of install_colordiff method:

def install_colordiff():
    # check if colordiff is installed
    out = subprocess.check_output(['apt-cache', 'policy', 'colordiff'])

    if re.search(r'Installed:\s*\(none\)', str(out)):
        # install with "sudo" if not exists
        subprocess.call(['sudo', 'apt-get', 'install', 'colordiff'])
        print('colordiff has been installed')

Step 4: Add useful Git aliases

We’re almost finished. Last step is to add a few useful Git aliases. For this, we will need to read content from our gitconfig.dist and append it in your ~/.gitconfig file.

def add_git_aliases():
    # make absolute path for new .gitconfig in user home directory
    home_gitconfig_path = os.path.join(home_dir, '.gitconfig')
    if not os.path.exists(home_gitconfig_path):
        # create .gitconfig file if it doesn't exists
        subprocess.call(['cp', gitconfig_dist_filename, home_gitconfig_path])
        print('.gitconfig file has been created')
        return

In this example, I assume I do not have ~/.gitconfig or have it empty. You might want to perform additional logic to merge ditconfig.dist with your existing config. In this case we will use Python module called configparser (configuration file parser). Let’s import this module to our script and modify add_git_aliases() function:

# in case of using Python2 use ConfigParser instead of configparser
import configparser

def add_git_aliases():
    # make absolute path for new .gitconfig in user home directory
    home_gitconfig_path = os.path.join(home_dir, '.gitconfig')
    if not os.path.exists(home_gitconfig_path):
        # create .gitconfig file if it doesn't exists
        subprocess.call(['cp', gitconfig_dist_filename, home_gitconfig_path])
        print('.gitconfig file has been created')
        return

    # optional (if exists, merge 2 configs and write the result)
    config = configparser.ConfigParser()
    config.read(home_gitconfig_path)
    config.read(gitconfig_dist_filename)

    # set up config file name
    config.filename = '.gitconfig'

    # write merged config to disk
    file_handler = open(home_gitconfig_path, 'w')
    config.write(file_handler)
    file_handler.close()

    print('.gitconfig file has been updated\n')

Step 5: Complete script

Now we have all the methods implemented we need to finish the script. Add those methods in a sequence execution and it is complete.

# execute changes
install_colordiff()
add_git_aliases()
bashrc_content = show_branch_in_shell(
    color_prompt_enable(bashrc_content)
)

Basic implementation is now finished. Feel free to add your own improvements to the script, such as more Linux distributives support, version control, custom useful aliases, and share this experience with other users!

With you newly gained knowledge and experience in Python you can now create your own Linux system administration scripts. Let me know what you came up with, please share them using the comments.

Source references:

  1. Script source on GitHub
  2. Python operating system interfaces module documentation
  3. Python regular expressions module documentation
  4. Python subprocess module documentation
  5. Python configparser module documentation

Leave a Reply

Your email address will not be published. Required fields are marked *