simpler way to have python2 compatibility

examples/parallel/pi/pidigits.py

@@ -1,161 +1,162 b''
 """Compute statistics on the digits of pi.
 
 This uses precomputed digits of pi from the website
 of Professor Yasumasa Kanada at the University of
 Tokoyo: http://www.super-computing.org/
 
 Currently, there are only functions to read the
 .txt (non-compressed, non-binary) files, but adding
 support for compression and binary files would be
 straightforward.
 
 This focuses on computing the number of times that
 all 1, 2, n digits sequences occur in the digits of pi.
 If the digits of pi are truly random, these frequencies
 should be equal.
 """
 
 # Import statements
 from __future__ import division, with_statement
 
 import numpy as np
 from matplotlib import pyplot as plt
 
-# Top-level functions
+try : #python2
+    from urllib import urlretrieve
+except : #python3
+    from urllib.request import urlretrieve
+
+	# Top-level functions
 
 def fetch_pi_file(filename):
     """This will download a segment of pi from super-computing.org
     if the file is not already present.
     """
     import os, urllib
     ftpdir="ftp://pi.super-computing.org/.2/pi200m/"
     if os.path.exists(filename):
         # we already have it
         return
     else:
         # download it
-        try : #python2
-            urllib.urlretrieve(ftpdir+filename,filename)
-        except : #python3
-            import urllib.request
-            urllib.request.urlretrieve(ftpdir+filename,filename)
+        urlretrieve(ftpdir+filename,filename)
 
 def compute_one_digit_freqs(filename):
     """
     Read digits of pi from a file and compute the 1 digit frequencies.
     """
     d = txt_file_to_digits(filename)
     freqs = one_digit_freqs(d)
     return freqs
 
 def compute_two_digit_freqs(filename):
     """
     Read digits of pi from a file and compute the 2 digit frequencies.
     """
     d = txt_file_to_digits(filename)
     freqs = two_digit_freqs(d)
     return freqs
 
 def reduce_freqs(freqlist):
     """
     Add up a list of freq counts to get the total counts.
     """
     allfreqs = np.zeros_like(freqlist[0])
     for f in freqlist:
         allfreqs += f
     return allfreqs
 
 def compute_n_digit_freqs(filename, n):
     """
     Read digits of pi from a file and compute the n digit frequencies.
     """
     d = txt_file_to_digits(filename)
     freqs = n_digit_freqs(d, n)
     return freqs
 
 # Read digits from a txt file
 
 def txt_file_to_digits(filename, the_type=str):
     """
     Yield the digits of pi read from a .txt file.
     """
     with open(filename, 'r') as f:
         for line in f.readlines():
             for c in line:
                 if c != '\n' and c!= ' ':
                     yield the_type(c)
 
 # Actual counting functions
 
 def one_digit_freqs(digits, normalize=False):
     """
     Consume digits of pi and compute 1 digit freq. counts.
     """
     freqs = np.zeros(10, dtype='i4')
     for d in digits:
         freqs[int(d)] += 1
     if normalize:
         freqs = freqs/freqs.sum()
     return freqs
 
 def two_digit_freqs(digits, normalize=False):
     """
     Consume digits of pi and compute 2 digits freq. counts.
     """
     freqs = np.zeros(100, dtype='i4')
     last = next(digits)
     this = next(digits)
     for d in digits:
         index = int(last + this)
         freqs[index] += 1
         last = this
         this = d
     if normalize:
         freqs = freqs/freqs.sum()
     return freqs
 
 def n_digit_freqs(digits, n, normalize=False):
     """
     Consume digits of pi and compute n digits freq. counts.
 
     This should only be used for 1-6 digits.
     """
     freqs = np.zeros(pow(10,n), dtype='i4')
     current = np.zeros(n, dtype=int)
     for i in range(n):
         current[i] = next(digits)
     for d in digits:
         index = int(''.join(map(str, current)))
         freqs[index] += 1
         current[0:-1] = current[1:]
         current[-1] = d
     if normalize:
         freqs = freqs/freqs.sum()
     return freqs
 
 # Plotting functions
 
 def plot_two_digit_freqs(f2):
     """
     Plot two digits frequency counts using matplotlib.
     """
     f2_copy = f2.copy()
     f2_copy.shape = (10,10)
     ax = plt.matshow(f2_copy)
     plt.colorbar()
     for i in range(10):
         for j in range(10):
             plt.text(i-0.2, j+0.2, str(j)+str(i))
     plt.ylabel('First digit')
     plt.xlabel('Second digit')
     return ax
 
 def plot_one_digit_freqs(f1):
     """
     Plot one digit frequency counts using matplotlib.
     """
     ax = plt.plot(f1,'bo-')
     plt.title('Single digit counts in pi')
     plt.xlabel('Digit')
     plt.ylabel('Count')
     return ax