Python tricks

List of code snippets I use all the time!

Pandas

Import pandas and numpy:

import pandas as pd
import numpy as np

Make random dataframe with three columns:

df = pd.DataFrame(np.random.rand(10,3), columns=list('ABC'))

Select:

# Boolean indexing (remember the parentheses)
df[(df.A < 0.5) & (df.B > 0.5)]
 
# Alternative, using query which depends on numexpr
df.query('A < 0.5 & B > 0.5')

Project:

# One columns
df.A
 
# Multiple columns
# there may be another shorter way, but I don't know it
df.loc[:,list('AB')]

Often used snippets

Dates

Difference (in days) between two dates:

from datetime import date
 
d1 = date(2013,1,1)
d2 = date(2013,9,13)
abs(d2-d1).days

directory-of-script snippet

os.path.dirname(os.path.realpath(__file__))
# combine with
os.path.join(os.path.dirname(os.path.realpath(__file__)), 'foo','bar','baz.txt')

PostgreSQL-connect-query snippet

import psycopg2
 
conn = psycopg2.connect("host='localhost' user='xxx'  password='yyy' dbname='zzz'")
cur = conn.cursor()
cur.execute("""SELECT * from foo;""")
rows = cur.fetchall()
for row in rows:
  print "   ", row[0]
conn.close()

Input parsing functions

Expand input-file args:

# input_data: e.g. 'file.txt' or '*.txt' or 'foo/file.txt' 'bar/file.txt'
filenames = [glob.glob(pathexpr) for pathexpr in input_data]
filenames = [item for sublist in filenames for item in sublist]

Parse key-value pair strings like ‘x=42.0,y=1’:

kvp = lambda elem,t,i: t(elem.split('=')[i])
parse_kvp_str = lambda args : dict([(kvp(elem,str,0), kvp(elem,float,1)) for elem in args.split(',')])
parse_kvp_str('x=42.0,y=1')

Postgres database functions

Upper case in Python (just for example):

-- create extension plpythonu;
 
CREATE OR REPLACE FUNCTION python_upper
(
  input text
) RETURNS text AS
$$
  return input.upper()
$$ LANGUAGE plpythonu STRICT;

Convert IP address from text to integer:

CREATE FUNCTION ip2int(input text) RETURNS integer
LANGUAGE plpythonu
AS $$
if 'struct' in SD:
    struct = SD['struct']
else:
    import struct
    SD['struct'] = struct
 
if 'socket' in SD:
    socket = SD['socket']
else:
    import socket
    SD['socket'] = socket
 
return struct.unpack("!I", socket.inet_aton(input))[0]
$$;

Convert IP address from integer to text:

CREATE FUNCTION int2ip(input integer) RETURNS text
LANGUAGE plpythonu
AS $$
if 'struct' in SD:
    struct = SD['struct']
else:
    import struct
    SD['struct'] = struct
 
if 'socket' in SD:
    socket = SD['socket']
else:
    import socket
    SD['socket'] = socket
 
return socket.inet_ntoa(struct.pack("!I", input))
$$;

Commandline options

optparse-commandline-options snippet

from optparse import OptionParser
 
usage = "usage: %prog [options] arg "
parser = OptionParser(usage=usage)
parser.add_option("-x", "--some-option-x", dest="x", default=42.0, type="float",
                  help="a floating point option")
(options, args) = parser.parse_args()
print options.x
print args[0]

print-in-place (progress bar) snippet

import time
import sys
for progress in range(100):
  time.sleep(0.1)
  sys.stdout.write("Download progress: %d%%   \r" % (progress) ) 
  sys.stdout.flush()

Packaging snippets

poor-mans-python-executable trick

Learned this trick from voidspace. The trick uses two files (__main__.py and hashbang.txt):

__main__.py:

print 'Hello world'

hashbang.txt (adding a newline after ‘python2.6’ is important):

#!/usr/bin/env python2.6

Build an “executable”:

zip main.zip __main__.py
cat hashbang.txt main.zip > hello
rm main.zip
chmod u+x hello

Run “executable”:

$ ./hello
Hello world

import-class-from-file trick

Import class MyClass from a module file (adapted from stackoverflow):

import imp
mod = imp.load_source('name.of.module', 'path/to/module.py')
obj = mod.MyClass()

Occusional-usage snippets

Extract words from string

words = lambda text: ''.join(c if c.isalnum() else ' ' for c in text).split()
words('Johnny.Appleseed!is:a*good&farmer')
# ['Johnny', 'Appleseed', 'is', 'a', 'good', 'farmer']

IP address to integer and back

import struct
import socket
 
def ip2int(addr):                                                               
    return struct.unpack("!I", socket.inet_aton(addr))[0]                       
 
 
def int2ip(addr):                                                               
    return socket.inet_ntoa(struct.pack("!I", addr))

Fluent Python Interface

Copied from riaanvddool.

# Fluent Interface Definition
class sql:
  class select:
    def __init__(self, dbcolumn, context=None):
      self.dbcolumn = dbcolumn
      self.context = context
    def select(self, dbcolumn):
      return self.__class__(dbcolumn,self)
 
# Demo
q = sql.select('foo').select('bar')
print q.dbcolumn  #bar
print q.context.dbcolumn  #foo

Flatten a nested lists

def flatten(elems):
	"""
	[['a'], ['b','c',['d'],'e',['f','g']]]
	"""
	stack = [elems]
	top = stack.pop()
	while top:		
		head, tail = top[0], top[1:]
		if tail: stack.append(tail)
		if not isinstance(head, list): yield head			
		else: stack.append(head)
		if stack: top = stack.pop()
		else: break

snap rounding

EPSILON = 0.000001
snap_ceil = lambda x: math.ceil(x) if abs(x - round(x)) > EPSILON else round(x)
snap_floor = lambda x: math.floor(x) if abs(x - round(x)) > EPSILON else round(x)

merge-two-dictionaries snippet

x = {'a': 42}
y = {'b': 127}
z = dict(x.items() + y.items())
# z = {'a': 42, 'b': 127}

anonymous-object snippet

Adapted from stackoverflow:

class Anon(object):
    def __new__(cls, **attrs):
        result = object.__new__(cls)
        result.__dict__ = attrs
        return result

Alternative:

class Anon(object):
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
 
    def __repr__(self):
        return self.__str__()
 
    def __str__(self):
        return ", ".join(["%s=%s" % (key,value) for key,value in self.__dict__.items()])

geocoding snippet

Geocoding an address with geopy.geocoders (wraps Bing, Yahoo, Google, GeoNames etc.)

from geopy import geocoders
geocoder = geocoders.GoogleV3()
geocoder.geocode( 'universitetsparken 5, 2100 københavn', exactly_one=False )
# [(u'Universitetsparken 5, University of Copenhagen, 2100 Copenhagen, Denmark', 
# (55.7013251, 12.5608664))]

generate-random-word snippet

Function that returns a random word (could also use random.choice with this list of words):

import string, random
randword = lambda n: "".join([random.choice(string.letters) for i in range(n)])

setdefault tricks

Increment (and initialize) value:

d = {}
d[2] = d.setdefault(2,39) + 1
d[2] = d.setdefault(2,39) + 1
d[2] = d.setdefault(2,39) + 1
d[2] # value is 42

Append value to (possibly uninitialized) list stored under a key in dictionary:

d = {}
d.setdefault(2, []).append(42)
d.setdefault(2, []).append(127)
d[2] # value is [42, 127]

Binary tricks

add-integers-using-XOR snippet

Swap two integer variables using the XOR swap algorithm:

x = 42
y = 127
x = x ^ y
y = y ^ x
x = x ^ y
x # value is 127
y # value is 42

Leave a Reply