#!/usr/bin/env python
# -*- coding: utf-8 -*-
 
# Takes a list of directories and prints the top n lines of each .txt file found
# Known bugs:
#   - Files without line breaks won't have their content listed.
 
import sys
import os
 
def getFirstNLines(filename, linesToRead):
  firstLines=""
  try:
    with open(filename,'r') as fileHandle:
        try:
          firstLines = [next(fileHandle) for x in range(linesToRead)]
        except StopIteration:
          pass #not enough lines in the file. Ignore and move on.
  except IOError:
    print("Unable to read file: " + filename)
  return firstLines
 
def getFilesInDir(path, ext):
  files = []
  for file in os.listdir(path):
    if file.endswith(ext):
        files.append(path+file)
  return files
 
def getFilesInDirectories(paths, extension):
    files = []
    for dir in paths:
        try:
            files.extend(getFilesInDir(dir, extension));
        except IOError:
            print("Unable to read path: " + dir)
        except OSError:
            print("No such file or directory: " + dir)
    return files
 
def prettyPrint(filePath, content):
  title = os.path.basename(filePath) #print only filename. might be a misstake -  user probably wants to  know where the file was located too?
  print(title + ":")
  for line in content:
    print("\t" + line.rstrip()); #rstrip to rid of trailing newlines
 
def main():
  if len(sys.argv) == 1: #1st argument is always the script name
    print("Not enough arguments provided")
    return
  FILE_EXT = ".txt"
  LINES_TO_READ = 2
  files = getFilesInDirectories(sys.argv[1:], FILE_EXT)  #slice the array to ignore 1st argument 
  for filePath in files:
    prettyPrint(filePath, getFirstNLines(filePath, LINES_TO_READ))
 
if __name__ == "__main__":
    main()