<commit_before>#!/usr/bin/env python

__version__ = "0.6"

import base64
import codecs
import locale
import os
import re
import shutil
import stat
import sys

from datetime import datetime
from optparse import OptionParser
from xml.dom.pulldom import START_ELEMENT, END_ELEMENT, parse
from xml.dom.minidom import Node

try:
    import pyexiv2
except ImportError:
    pyexiv2 = None

# To allow Unicode characters to be displayed
# (see http://wiki.python.org/moin/PrintFails)
sys.stdout = codecs.getwriter(locale.getpreferredencoding())(sys.stdout)
sys.stderr = codecs.getwriter(locale.getpreferredencoding())(sys.stderr)

class iPhotoLibraryError(Exception):
    pass

class iPhotoLibrary(object):
    def __init__(self, albumDir, destDir, use_album=False, use_date=False,
                 use_faces=False, use_metadata=False, quiet=False):
        self.use_album = use_album
        self.use_date =  use_date
        self.use_faces = use_faces
        self.use_metadata = use_metadata
        self.dest_dir = destDir
        self.quiet = quiet
        self.albums = []
        self.keywords = {}
        self.faces = {}
        self.images = {}
        albumDataXml = os.path.join(albumDir, "AlbumData.xml")
        self.status("* Parsing iPhoto Library data... ")
        self.parseAlbumData(albumDataXml)
        self.status("Done.\n")

    major_version = 2
    minor_version = 0
    interesting_image_keys = [
        'ImagePath', 'Rating', 'Keywords', 'Caption', 'Comment', 'Faces',
        'face key'
    ]
    apple_epoch = 978307200

    def parseAlbumData(self, filename):
        """
        Parse an iPhoto AlbumData.xml file, keeping the interesting
        bits.
        """
        doc = parse(filename)
        stack = []
        last_top_key = None
        if self.use_album:
            album_list_key = "List of Albums"
        else:
            album_list_key = "List of Rolls"
        for event, node in doc:
            if event == START_ELEMENT:
                stack.append(node)
                level = len(stack)
                if level == 3:
                    if node.nodeName == 'key':
                        doc.expandNode(node)
                        last_top_key = self.getText(node)
                        stack.pop()
                    elif last_top_key == 'List of Keywords':
                        doc.expandNode(node)
                        self.keywords = self.dePlist(node)
                        stack.pop()
                    elif last_top_key == 'List of Faces':
                        doc.expandNode(node)
                        self.faces = dict([
                            (k, v['name']) for k,v in
                             self.dePlist(node, ['name']).items()
                        ])
                        stack.pop()
                    elif last_top_key == 'Major Version':
                        doc.expandNode(node)
                        major_version = self.dePlist(node)
                        stack.pop()
                        if major_version != self.major_version:
                            raise iPhotoLibraryError, \
                            "Sorry, I can't understand version %i iPhoto Libraries." % major_version
                    elif last_top_key == 'Minor Version':
                        doc.expandNode(node)
                        minor_version = self.dePlist(node)
                        stack.pop()
                        if minor_version > self.minor_version:
                            self.status(
                                "\nI don't recognise iPhoto libraries when the minor version is %i, but let's try anyway.\n" % minor_version,
                                force=True
                            )

                elif level == 4:
                    # process large items individually so we don't
                    # load them all into memory.
                    if last_top_key == album_list_key:
                        doc.expandNode(node)
                        self.albums.append(self.dePlist(node))
                        stack.pop()
                    elif last_top_key == 'Master Image List':
                        doc.expandNode(node)
                        if node.nodeName == 'key':
                            last_image_key = self.getText(node)
                        else:
                            self.images[last_image_key] = self.dePlist(
                                node, self.interesting_image_keys
                            )
                        stack.pop()
            elif event == END_ELEMENT:
                stack.pop()

    def dePlist(self, node, interesting_keys=None):
        """
        Given a DOM node, convert the plist (fragment) it refers to and
        return the corresponding Python data structure.

        If interesting_keys is a list, "dict" keys will be filtered so that
        only those nominated are returned (for ALL descendant dicts). Numeric
        keys aren't filtered.
        """
        ik = interesting_keys
        dtype = node.nodeName
        if dtype == 'string':
            return self.getText(node)
        elif dtype == 'integer':
            try:
                return int(self.getText(node))
            except ValueError:
                raise iPhotoLibraryError, \
                "Corrupted Library; unexpected value '%s' for integer" % \
                    self.getText(node)
        elif dtype == 'real':
            try:
                return float(self.getText(node))
            except ValueError:
                raise iPhotoLibraryError, \
                "Corrupted Library; unexpected value '%s' for real" % \
                    self.getText(node)
        elif dtype == 'array':
            return [self.dePlist(c, ik) for c in node.childNodes \
                    if c.nodeType == Node.ELEMENT_NODE]
        elif dtype == 'dict':
            d = {}
            last_key = None
            for c in node.childNodes:
                if c.nodeType != Node.ELEMENT_NODE:
                    continue
                # TODO: catch out-of-order keys/values
                if c.nodeName == 'key':
                    last_key = self.getText(c)
                else: # value
                    if interesting_keys: # check to see if we're interested
                        if last_key not in interesting_keys \
                          and not last_key.isdigit():
                            continue # nope.
                    d[intern(str(last_key))] = self.dePlist(c, ik)
            return d
        elif dtype == 'true':
            return True
        elif dtype == 'false':
            return False
        elif dtype == 'data':
            return base64.decodestring(self.getText(c))
        elif dtype == 'date':
            return self.appleDate(self.getText(c))
        else:
            raise Exception, "Don't know what a %s is." % dtype

    @staticmethod
    def getText(element, default=None):
        if element is None:
            return default
        if len(element.childNodes) == 0:
            return None
        else:
            return "".join([n.nodeValue for n in element.childNodes])

    def walk(self, funcs):
        """
        Walk through the events or albums (depending on the value of albums)
        in this library and apply each function in the list funcs to each
        image, calling it as:
           func(folderName, folderDate, imageId)
        where:
         - folderName is the name the folder,
         - folderDate is the date of the folder, and
         - imageId is the string identifier for the image.
        """
        if self.use_album:
            targetName = "AlbumName"
            albums = [a for a in self.albums if
                      a.get("Album Type", None) == "Regular"]
        else:
            targetName = "RollName"
            albums = self.albums
        i = 0
        for folder in albums:
            i += 1
            folderName = folder[targetName]
            if self.use_album:
                folderDate = None
            else:
                folderDate = self.appleDate(folder["RollDateAsTimerInterval"])
            images = folder["KeyList"]
            self.status("* Processing %i of %i: %s (%i images)...\n" % (
                i,
                len(albums),
                folderName,
                len(images)
            ))
            for imageId in images:
                for func in funcs:
                    func(imageId, folderName, folderDate)
            self.status("\n")

    def copyImage(self, imageId, folderName, folderDate):
        """
        Copy an image from the library to a folder in the dest_dir. The
        name of the folder is based on folderName and folderDate; if
        folderDate is None, it's only based upon the folderName.

        If use_metadata is True, also write the image metadata from the library
        to the copy. If use_faces is True, faces will be saved as keywords.
        """
        try:
            image = self.images[imageId]
        except KeyError:
            raise iPhotoLibraryError, "Can't find image #%s" % imageId

        if folderDate and self.use_date:
            date = '%(year)d-%(month)02d-%(day)02d' % {
                'year': folderDate.year,
                'month': folderDate.month,
                'day': folderDate.day
            }
            if re.match("[A-Z][a-z]{2} [0-9]{1,2}, [0-9]{4}", folderName):
                outputPath = date
            else:
                outputPath = date + " " + folderName
        else:
            outputPath = folderName
        targetFileDir = os.path.join(self.dest_dir, outputPath)

        if not os.path.exists(targetFileDir):
            try:
                os.makedirs(targetFileDir)
            except OSError, why:
                raise iPhotoLibraryError, \
                    "Can't create %s: %s" % (targetFileDir, why[1])
            self.status("  Created %s\n" % targetFileDir)

        mFilePath = image["ImagePath"]
        basename = os.path.basename(mFilePath)
        tFilePath = os.path.join(targetFileDir, basename)

        # Skip unchanged files, unless we're writing metadata.
        if not self.use_metadata and os.path.exists(tFilePath):
            mStat = os.stat(mFilePath)
            tStat = os.stat(tFilePath)
            if abs(tStat[stat.ST_MTIME] - mStat[stat.ST_MTIME]) <= 10 or \
              tStat[stat.ST_SIZE] == mStat[stat.ST_SIZE]:
                self.status("-")
                return

        shutil.copy2(mFilePath, tFilePath)
        md_written = False
        if self.use_metadata:
            md_written = self.writePhotoMD(imageId, tFilePath)
        if md_written:
            self.status("+")
        else:
            self.status(".")

    def writePhotoMD(self, imageId, filePath=None):
        """
        Write the metadata from the library for imageId to filePath.
        If filePath is None, write it to the photo in the library.
        If use_faces is True, iPhoto face names will be written to
        keywords.
        """
        try:
            image = self.images[imageId]
        except KeyError:
            raise iPhotoLibraryError, "Can't find image #%s" % imageId
        if not filePath:
            filePath = image['ImagePath']

        caption = image.get("Caption", None)
        rating = image.get("Rating", None)
        comment = image.get("Comment", None)
        keywords = set([self.keywords[k] for k in image.get("Keywords", [])])
        if self.use_faces:
            keywords.update([self.faces[f['face key']]
                             for f in image.get("Faces", [])
                             if self.faces.has_key(f['face key'])]
            )

        if caption or comment or rating or keywords:
            try:
                md = pyexiv2.ImageMetadata(filePath)
                md.read()
                if caption:
                    md["Iptc.Application2.Headline"] = [caption]
                if rating:
                    md["Xmp.xmp.Rating"] = rating
                if comment:
                    md["Iptc.Application2.Caption"] = [comment]
                if keywords:
                    md["Iptc.Application2.Keywords"] = list(keywords)
                md.write(preserve_timestamps=True)
                return True
            except IOError, why:
                self.status("\nProblem setting metadata (%s) on %s\n" % (
                    why, filePath
                ))
        return False

    def appleDate(self, text):
        try:
            return datetime.utcfromtimestamp(self.apple_epoch + float(text))
        except (ValueError, TypeError):
            raise iPhotoLibraryError, \
            "Corrupted Library; unexpected value '%s' for date" % text

    def status(self, msg, force=False):
        if force or not self.quiet:
            sys.stdout.write(msg)
            sys.stdout.flush()

def error(msg):
    sys.stderr.write("\n%s\n" % msg)
    sys.exit(1)

if __name__ == '__main__':
    usage   = "Usage: %prog [options] <iPhoto Library dir> <destination dir>"
    version = "exportiphoto version %s" % __version__
    option_parser = OptionParser(usage=usage, version=version)
    option_parser.set_defaults(
        test=False,
        albums=False,
        metadata=False,
        faces=False,
        quiet=False,
        date=True
    )

    option_parser.add_option("-a", "--albums",
                             action="store_true", dest="albums",
                             help="use albums instead of events"
    )

    option_parser.add_option("-q", "--quiet",
                             action="store_true", dest="quiet",
                             help="use quiet mode"
    )

    option_parser.add_option("-d", "--date",
                             action="store_false", dest="date",
                             help="stop use date prefix in folder name"
    )

    if pyexiv2:
        option_parser.add_option("-m", "--metadata",
                                 action="store_true", dest="metadata",
                                 help="write metadata to images"
        )

        option_parser.add_option("-f", "--faces",
                                 action="store_true", dest="faces",
                                 help="store faces as keywords (requires -m)"
        )

    (options, args) = option_parser.parse_args()

    if len(args) != 2:
        option_parser.error(
            "Please specify an iPhoto library and a destination."
        )

    try:
        library = iPhotoLibrary(args[0], # src
                                args[1], # dest
                                use_album=options.albums,
                                use_date=options.date,
                                use_faces=options.faces,
                                use_metadata=options.metadata,
                                quiet=options.quiet)
        def copyImage(imageId, folderName, folderDate):
            library.copyImage(imageId, folderName, folderDate)
    except iPhotoLibraryError, why:
        error(why[0])
    except KeyboardInterrupt:
        error("Interrupted.")
    try:
        library.walk([copyImage])
    except iPhotoLibraryError, why:
        error(why[0])
    except KeyboardInterrupt:
        error("Interrupted. Copy may be incomplete.")
<commit_msg>Added deconfliction of output directories.<commit_after>#!/usr/bin/env python

__version__ = "0.6"

import base64
import codecs
import locale
import os
import re
import shutil
import stat
import sys

from datetime import datetime
from optparse import OptionParser
from xml.dom.pulldom import START_ELEMENT, END_ELEMENT, parse
from xml.dom.minidom import Node

try:
    import pyexiv2
except ImportError:
    pyexiv2 = None

# To allow Unicode characters to be displayed
# (see http://wiki.python.org/moin/PrintFails)
sys.stdout = codecs.getwriter(locale.getpreferredencoding())(sys.stdout)
sys.stderr = codecs.getwriter(locale.getpreferredencoding())(sys.stderr)

class iPhotoLibraryError(Exception):
    pass

class iPhotoLibrary(object):
    def __init__(self, albumDir, destDir, use_album=False, use_date=False,
                 use_faces=False, use_metadata=False, quiet=False):
        self.use_album = use_album
        self.use_date =  use_date
        self.use_faces = use_faces
        self.use_metadata = use_metadata
        self.dest_dir = destDir
        self.output_dirs = set()
        self.quiet = quiet
        self.albums = []
        self.keywords = {}
        self.faces = {}
        self.images = {}
        albumDataXml = os.path.join(albumDir, "AlbumData.xml")
        self.status("* Parsing iPhoto Library data... ")
        self.parseAlbumData(albumDataXml)
        self.status("Done.\n")

    major_version = 2
    minor_version = 0
    interesting_image_keys = [
        'ImagePath', 'Rating', 'Keywords', 'Caption', 'Comment', 'Faces',
        'face key'
    ]
    apple_epoch = 978307200

    def parseAlbumData(self, filename):
        """
        Parse an iPhoto AlbumData.xml file, keeping the interesting
        bits.
        """
        doc = parse(filename)
        stack = []
        last_top_key = None
        if self.use_album:
            album_list_key = "List of Albums"
        else:
            album_list_key = "List of Rolls"
        for event, node in doc:
            if event == START_ELEMENT:
                stack.append(node)
                level = len(stack)
                if level == 3:
                    if node.nodeName == 'key':
                        doc.expandNode(node)
                        last_top_key = self.getText(node)
                        stack.pop()
                    elif last_top_key == 'List of Keywords':
                        doc.expandNode(node)
                        self.keywords = self.dePlist(node)
                        stack.pop()
                    elif last_top_key == 'List of Faces':
                        doc.expandNode(node)
                        self.faces = dict([
                            (k, v['name']) for k,v in
                             self.dePlist(node, ['name']).items()
                        ])
                        stack.pop()
                    elif last_top_key == 'Major Version':
                        doc.expandNode(node)
                        major_version = self.dePlist(node)
                        stack.pop()
                        if major_version != self.major_version:
                            raise iPhotoLibraryError, \
                            "Sorry, I can't understand version %i iPhoto Libraries." % major_version
                    elif last_top_key == 'Minor Version':
                        doc.expandNode(node)
                        minor_version = self.dePlist(node)
                        stack.pop()
                        if minor_version > self.minor_version:
                            self.status(
                                "\nI don't recognise iPhoto libraries when the minor version is %i, but let's try anyway.\n" % minor_version,
                                force=True
                            )

                elif level == 4:
                    # process large items individually so we don't
                    # load them all into memory.
                    if last_top_key == album_list_key:
                        doc.expandNode(node)
                        self.albums.append(self.dePlist(node))
                        stack.pop()
                    elif last_top_key == 'Master Image List':
                        doc.expandNode(node)
                        if node.nodeName == 'key':
                            last_image_key = self.getText(node)
                        else:
                            self.images[last_image_key] = self.dePlist(
                                node, self.interesting_image_keys
                            )
                        stack.pop()
            elif event == END_ELEMENT:
                stack.pop()

    def dePlist(self, node, interesting_keys=None):
        """
        Given a DOM node, convert the plist (fragment) it refers to and
        return the corresponding Python data structure.

        If interesting_keys is a list, "dict" keys will be filtered so that
        only those nominated are returned (for ALL descendant dicts). Numeric
        keys aren't filtered.
        """
        ik = interesting_keys
        dtype = node.nodeName
        if dtype == 'string':
            return self.getText(node)
        elif dtype == 'integer':
            try:
                return int(self.getText(node))
            except ValueError:
                raise iPhotoLibraryError, \
                "Corrupted Library; unexpected value '%s' for integer" % \
                    self.getText(node)
        elif dtype == 'real':
            try:
                return float(self.getText(node))
            except ValueError:
                raise iPhotoLibraryError, \
                "Corrupted Library; unexpected value '%s' for real" % \
                    self.getText(node)
        elif dtype == 'array':
            return [self.dePlist(c, ik) for c in node.childNodes \
                    if c.nodeType == Node.ELEMENT_NODE]
        elif dtype == 'dict':
            d = {}
            last_key = None
            for c in node.childNodes:
                if c.nodeType != Node.ELEMENT_NODE:
                    continue
                # TODO: catch out-of-order keys/values
                if c.nodeName == 'key':
                    last_key = self.getText(c)
                else: # value
                    if interesting_keys: # check to see if we're interested
                        if last_key not in interesting_keys \
                          and not last_key.isdigit():
                            continue # nope.
                    d[intern(str(last_key))] = self.dePlist(c, ik)
            return d
        elif dtype == 'true':
            return True
        elif dtype == 'false':
            return False
        elif dtype == 'data':
            return base64.decodestring(self.getText(c))
        elif dtype == 'date':
            return self.appleDate(self.getText(c))
        else:
            raise Exception, "Don't know what a %s is." % dtype

    @staticmethod
    def getText(element, default=None):
        if element is None:
            return default
        if len(element.childNodes) == 0:
            return None
        else:
            return "".join([n.nodeValue for n in element.childNodes])

    def walk(self, funcs):
        """
        Walk through the events or albums (depending on the value of albums)
        in this library and apply each function in the list funcs to each
        image, calling it as:
           func(folderName, folderDate, imageId)
        where:
         - folderName is the name the folder,
         - folderDate is the date of the folder, and
         - imageId is the string identifier for the image.
        """
        if self.use_album:
            targetName = "AlbumName"
            albums = [a for a in self.albums if
                      a.get("Album Type", None) == "Regular"]
        else:
            targetName = "RollName"
            albums = self.albums
        i = 0
        for folder in albums:
            i += 1
            if self.use_album:
                folderDate = None
            else:
                folderDate = self.appleDate(folder["RollDateAsTimerInterval"])
            images = folder["KeyList"]

            folderName = folder[targetName]
            if folderDate and self.use_date:
                date = '%(year)d-%(month)02d-%(day)02d' % {
                    'year': folderDate.year,
                    'month': folderDate.month,
                    'day': folderDate.day
                }
                if re.match("[A-Z][a-z]{2} [0-9]{1,2}, [0-9]{4}", folderName):
                    outputPath = date
                else:
                    outputPath = date + " " + folderName
            else:
                outputPath = folderName

            # Deconflict output directories
            targetFileDir = os.path.join(self.dest_dir, outputPath)
            j = 1
            while targetFileDir in self.output_dirs:
                targetFileDir = os.path.join(self.dest_dir, outputPath + " %02d"%j)
                j += 1
            self.output_dirs.add(targetFileDir)

            self.status("* Processing %i of %i: %s (%i images)...\n" % (
                i,
                len(albums),
                folderName,
                len(images)
            ))
            for imageId in images:
                for func in funcs:
                    func(imageId, targetFileDir, folderDate)
            self.status("\n")

    def copyImage(self, imageId, folderName, folderDate):
        """
        Copy an image from the library to a folder in the dest_dir. The
        name of the folder is based on folderName and folderDate; if
        folderDate is None, it's only based upon the folderName.

        If use_metadata is True, also write the image metadata from the library
        to the copy. If use_faces is True, faces will be saved as keywords.
        """
        try:
            image = self.images[imageId]
        except KeyError:
            raise iPhotoLibraryError, "Can't find image #%s" % imageId

        if not os.path.exists(folderName):
            try:
                os.makedirs(folderName)
            except OSError, why:
                raise iPhotoLibraryError, \
                    "Can't create %s: %s" % (folderName, why[1])
            self.status("  Created %s\n" % folderName)

        mFilePath = image["ImagePath"]
        basename = os.path.basename(mFilePath)
        tFilePath = os.path.join(folderName, basename)

        # Skip unchanged files, unless we're writing metadata.
        if not self.use_metadata and os.path.exists(tFilePath):
            mStat = os.stat(mFilePath)
            tStat = os.stat(tFilePath)
            if abs(tStat[stat.ST_MTIME] - mStat[stat.ST_MTIME]) <= 10 or \
              tStat[stat.ST_SIZE] == mStat[stat.ST_SIZE]:
                self.status("-")
                return

        shutil.copy2(mFilePath, tFilePath)
        md_written = False
        if self.use_metadata:
            md_written = self.writePhotoMD(imageId, tFilePath)
        if md_written:
            self.status("+")
        else:
            self.status(".")

    def writePhotoMD(self, imageId, filePath=None):
        """
        Write the metadata from the library for imageId to filePath.
        If filePath is None, write it to the photo in the library.
        If use_faces is True, iPhoto face names will be written to
        keywords.
        """
        try:
            image = self.images[imageId]
        except KeyError:
            raise iPhotoLibraryError, "Can't find image #%s" % imageId
        if not filePath:
            filePath = image['ImagePath']

        caption = image.get("Caption", None)
        rating = image.get("Rating", None)
        comment = image.get("Comment", None)
        keywords = set([self.keywords[k] for k in image.get("Keywords", [])])
        if self.use_faces:
            keywords.update([self.faces[f['face key']]
                             for f in image.get("Faces", [])
                             if self.faces.has_key(f['face key'])]
            )

        if caption or comment or rating or keywords:
            try:
                md = pyexiv2.ImageMetadata(filePath)
                md.read()
                if caption:
                    md["Iptc.Application2.Headline"] = [caption]
                if rating:
                    md["Xmp.xmp.Rating"] = rating
                if comment:
                    md["Iptc.Application2.Caption"] = [comment]
                if keywords:
                    md["Iptc.Application2.Keywords"] = list(keywords)
                md.write(preserve_timestamps=True)
                return True
            except IOError, why:
                self.status("\nProblem setting metadata (%s) on %s\n" % (
                    why, filePath
                ))
        return False

    def appleDate(self, text):
        try:
            return datetime.utcfromtimestamp(self.apple_epoch + float(text))
        except (ValueError, TypeError):
            raise iPhotoLibraryError, \
            "Corrupted Library; unexpected value '%s' for date" % text

    def status(self, msg, force=False):
        if force or not self.quiet:
            sys.stdout.write(msg)
            sys.stdout.flush()

def error(msg):
    sys.stderr.write("\n%s\n" % msg)
    sys.exit(1)

if __name__ == '__main__':
    usage   = "Usage: %prog [options] <iPhoto Library dir> <destination dir>"
    version = "exportiphoto version %s" % __version__
    option_parser = OptionParser(usage=usage, version=version)
    option_parser.set_defaults(
        test=False,
        albums=False,
        metadata=False,
        faces=False,
        quiet=False,
        date=True
    )

    option_parser.add_option("-a", "--albums",
                             action="store_true", dest="albums",
                             help="use albums instead of events"
    )

    option_parser.add_option("-q", "--quiet",
                             action="store_true", dest="quiet",
                             help="use quiet mode"
    )

    option_parser.add_option("-d", "--date",
                             action="store_false", dest="date",
                             help="stop use date prefix in folder name"
    )

    if pyexiv2:
        option_parser.add_option("-m", "--metadata",
                                 action="store_true", dest="metadata",
                                 help="write metadata to images"
        )

        option_parser.add_option("-f", "--faces",
                                 action="store_true", dest="faces",
                                 help="store faces as keywords (requires -m)"
        )

    (options, args) = option_parser.parse_args()

    if len(args) != 2:
        option_parser.error(
            "Please specify an iPhoto library and a destination."
        )

    try:
        library = iPhotoLibrary(args[0], # src
                                args[1], # dest
                                use_album=options.albums,
                                use_date=options.date,
                                use_faces=options.faces,
                                use_metadata=options.metadata,
                                quiet=options.quiet)
        def copyImage(imageId, folderName, folderDate):
            library.copyImage(imageId, folderName, folderDate)
    except iPhotoLibraryError, why:
        error(why[0])
    except KeyboardInterrupt:
        error("Interrupted.")
    try:
        library.walk([copyImage])
    except iPhotoLibraryError, why:
        error(why[0])
    except KeyboardInterrupt:
        error("Interrupted. Copy may be incomplete.")
<|endoftext|>


<commit_before>__version__ = '0.8'
<commit_msg>Version 0.9<commit_after>__version__ = '0.9'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Named-entity recognizer
"""
from pythainlp.corpus import download, get_file, stopwords
from pythainlp.tag import pos_tag
from pythainlp.tokenize import word_tokenize

try:
    import sklearn_crfsuite
except ImportError:
    from pythainlp.tools import install_package

    install_package("sklearn-crfsuite")
    try:
        import sklearn_crfsuite
    except ImportError:
        raise ImportError("ImportError: Try 'pip install sklearn-crfsuite'")

_WORD_TOKENIZER = "newmm"  # ตัวตัดคำ
_STOPWORDS = stopwords.words("thai")

def _is_thaichar(ch):  # เช็คว่าเป็น char ภาษาไทย
    ch_val = ord(ch)
    if ch_val >= 3584 and ch_val <= 3711:
        return True
    return False

def _is_thaiword(word):  # เช็คว่าเป็นคำภาษาไทย
    for ch in word:
        if ch != "." and not _is_thaichar(ch):
            return False
    return True

def _is_stopword(word):  # เช็คว่าเป็นคำฟุ่มเฟือย
    return word in _STOPWORDS

def _doc2features(doc, i):
    word = doc[i][0]
    postag = doc[i][1]
    # Features from current word
    features = {
        "word.word": word,
        "word.stopword": _is_stopword(word),
        "word.isthai": _is_thaiword(word),
        "word.isspace": word.isspace(),
        "postag": postag,
        "word.isdigit()": word.isdigit(),
    }
    if word.isdigit() and len(word) == 5:
        features["word.islen5"] = True
    if i > 0:
        prevword = doc[i - 1][0]
        postag1 = doc[i - 1][1]
        features["word.prevword"] = prevword
        features["word.previsspace"] = prevword.isspace()
        features["word.previsthai"] = _is_thaiword(prevword)
        features["word.prevstopword"] = _is_stopword(prevword)
        features["word.prepostag"] = postag1
        features["word.prevwordisdigit"] = prevword.isdigit()
    else:
        features["BOS"] = True  # Special "Beginning of Sequence" tag

    # Features from next word
    if i < len(doc) - 1:
        nextword = doc[i + 1][0]
        postag1 = doc[i + 1][1]
        features["word.nextword"] = nextword
        features["word.nextisspace"] = nextword.isspace()
        features["word.nextpostag"] = postag1
        features["word.nextisthai"] = _is_thaiword(nextword)
        features["word.nextstopword"] = _is_stopword(nextword)
        features["word.nextwordisdigit"] = nextword.isdigit()
    else:
        features["EOS"] = True  # Special "End of Sequence" tag

    return features

class ThaiNameRecognizer:
    def __init__(self):
        """
        Thai named-entity recognizer
        """
        self.__data_path = get_file("thainer")
        if not self.__data_path:
            download("thainer")
            self.__data_path = get_file("thainer")
        self.crf = sklearn_crfsuite.CRF(
            algorithm="lbfgs",
            c1=0.1,
            c2=0.1,
            max_iterations=500,
            all_possible_transitions=True,
            model_filename=self.__data_path,
        )

    def get_ner(self, text, pos=True):
        """
        Get named-entities in text

        :param string text: Thai text
        :param boolean pos: get Part-Of-Speech tag (True) or get not (False)

        :return: list of strings with name labels (and part-of-speech tags)

        **Example**::
            >>> from pythainlp.ner import ThaiNameRecognizer
            >>> ner=ThaiNameRecognizer()
            >>> ner.get_ner("วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา 14:49 น.")
            [('วันที่', 'JSBR', 'O'), (' ', 'NCMN', 'O'), ('15', 'NCNM', 'B-DATE'),
            (' ', 'NCMN', 'I-DATE'), ('ก.ย.', 'CMTR', 'I-DATE'), (' ', 'NCMN', 'I-DATE'),
            ('61', 'NCNM', 'I-DATE'), (' ', 'NCMN', 'O'), ('ทดสอบ', 'VACT', 'O'),
            ('ระบบ', 'NCMN', 'O'), ('เวลา', 'NCMN', 'O'), (' ', 'NCMN', 'O'),
            ('14', 'NCNM', 'B-TIME'), (':', 'PUNC', 'I-TIME'), ('49', 'NCNM', 'I-TIME'),
            (' ', 'NCMN', 'I-TIME'), ('น.', 'CMTR', 'I-TIME')]
            >>> ner.get_ner("วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา 14:49 น.",postag=False)
            [('วันที่', 'O'), (' ', 'O'), ('15', 'B-DATE'), (' ', 'I-DATE'),
            ('ก.ย.', 'I-DATE'), (' ', 'I-DATE'), ('61', 'I-DATE'), (' ', 'O'),
            ('ทดสอบ', 'O'), ('ระบบ', 'O'), ('เวลา', 'O'), (' ', 'O'), ('14', 'B-TIME'),
            (':', 'I-TIME'), ('49', 'I-TIME'), (' ', 'I-TIME'), ('น.', 'I-TIME')]
        """
        self.__tokens = word_tokenize(text, engine=_WORD_TOKENIZER)
        self.__pos_tags = pos_tag(self.__tokens, engine="perceptron")
        self.__x_test = self.__extract_features(
            [(data, self.__pos_tags[i][1]) for i, data in enumerate(self.__tokens)]
        )
        self.__y = self.crf.predict_single(self.__x_test)

        if pos:
            return [
                (self.__tokens[i], self.__pos_tags[i][1], data)
                for i, data in enumerate(self.__y)
            ]

        return [(self.__tokens[i], data) for i, data in enumerate(self.__y)]

    def __extract_features(self, doc):
        return [_doc2features(doc, i) for i in range(len(doc))]

    def __get_labels(self, doc):
        return [tag for (token, postag, tag) in doc]

    def __get_model(self):
        return self.crf
<commit_msg>Update NER doc<commit_after># -*- coding: utf-8 -*-
"""
Named-entity recognizer
"""
from pythainlp.corpus import download, get_file, stopwords
from pythainlp.tag import pos_tag
from pythainlp.tokenize import word_tokenize

try:
    import sklearn_crfsuite
except ImportError:
    from pythainlp.tools import install_package

    install_package("sklearn-crfsuite")
    try:
        import sklearn_crfsuite
    except ImportError:
        raise ImportError("ImportError: Try 'pip install sklearn-crfsuite'")

_WORD_TOKENIZER = "newmm"  # ตัวตัดคำ
_STOPWORDS = stopwords.words("thai")

def _is_thaichar(ch):  # เป็นอักษรไทยหรือไม่
    ch_val = ord(ch)
    if ch_val >= 3584 and ch_val <= 3711:
        return True
    return False

def _is_thaiword(word):  # เป็นคำที่มีแต่อักษรไทยหรือไม่
    for ch in word:
        if ch != "." and not _is_thaichar(ch):
            return False
    return True

def _is_stopword(word):  # เช็คว่าเป็นคำฟุ่มเฟือย
    return word in _STOPWORDS

def _doc2features(doc, i):
    word = doc[i][0]
    postag = doc[i][1]
    # Features from current word
    features = {
        "word.word": word,
        "word.stopword": _is_stopword(word),
        "word.isthai": _is_thaiword(word),
        "word.isspace": word.isspace(),
        "postag": postag,
        "word.isdigit()": word.isdigit(),
    }

    if word.isdigit() and len(word) == 5:
        features["word.islen5"] = True

    if i > 0:
        prevword = doc[i - 1][0]
        postag1 = doc[i - 1][1]
        features["word.prevword"] = prevword
        features["word.previsspace"] = prevword.isspace()
        features["word.previsthai"] = _is_thaiword(prevword)
        features["word.prevstopword"] = _is_stopword(prevword)
        features["word.prepostag"] = postag1
        features["word.prevwordisdigit"] = prevword.isdigit()
    else:
        features["BOS"] = True  # Special "Beginning of Sequence" tag

    # Features from next word
    if i < len(doc) - 1:
        nextword = doc[i + 1][0]
        postag1 = doc[i + 1][1]
        features["word.nextword"] = nextword
        features["word.nextisspace"] = nextword.isspace()
        features["word.nextpostag"] = postag1
        features["word.nextisthai"] = _is_thaiword(nextword)
        features["word.nextstopword"] = _is_stopword(nextword)
        features["word.nextwordisdigit"] = nextword.isdigit()
    else:
        features["EOS"] = True  # Special "End of Sequence" tag

    return features

class ThaiNameRecognizer:
    def __init__(self):
        """
        Thai named-entity recognizer
        """
        self.__data_path = get_file("thainer")
        if not self.__data_path:
            download("thainer")
            self.__data_path = get_file("thainer")
        self.crf = sklearn_crfsuite.CRF(
            algorithm="lbfgs",
            c1=0.1,
            c2=0.1,
            max_iterations=500,
            all_possible_transitions=True,
            model_filename=self.__data_path,
        )

    def get_ner(self, text, pos=True):
        """
        Get named-entities in text

        :param string text: Thai text
        :param boolean pos: get Part-Of-Speech tag (True) or get not (False)

        :return: list of strings with name labels (and part-of-speech tags)

        **Example**::
            >>> from pythainlp.ner import ThaiNameRecognizer
            >>> ner = ThaiNameRecognizer()
            >>> ner.get_ner("วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา 14:49 น.")
            [('วันที่', 'JSBR', 'O'), (' ', 'NCMN', 'O'), ('15', 'NCNM', 'B-DATE'),
            (' ', 'NCMN', 'I-DATE'), ('ก.ย.', 'CMTR', 'I-DATE'), (' ', 'NCMN', 'I-DATE'),
            ('61', 'NCNM', 'I-DATE'), (' ', 'NCMN', 'O'), ('ทดสอบ', 'VACT', 'O'),
            ('ระบบ', 'NCMN', 'O'), ('เวลา', 'NCMN', 'O'), (' ', 'NCMN', 'O'),
            ('14', 'NCNM', 'B-TIME'), (':', 'PUNC', 'I-TIME'), ('49', 'NCNM', 'I-TIME'),
            (' ', 'NCMN', 'I-TIME'), ('น.', 'CMTR', 'I-TIME')]
            >>> ner.get_ner("วันที่ 15 ก.ย. 61 ทดสอบระบบเวลา 14:49 น.", pos=False)
            [('วันที่', 'O'), (' ', 'O'), ('15', 'B-DATE'), (' ', 'I-DATE'),
            ('ก.ย.', 'I-DATE'), (' ', 'I-DATE'), ('61', 'I-DATE'), (' ', 'O'),
            ('ทดสอบ', 'O'), ('ระบบ', 'O'), ('เวลา', 'O'), (' ', 'O'), ('14', 'B-TIME'),
            (':', 'I-TIME'), ('49', 'I-TIME'), (' ', 'I-TIME'), ('น.', 'I-TIME')]
        """
        self.__tokens = word_tokenize(text, engine=_WORD_TOKENIZER)
        self.__pos_tags = pos_tag(self.__tokens, engine="perceptron")
        self.__x_test = self.__extract_features(
            [(data, self.__pos_tags[i][1]) for i, data in enumerate(self.__tokens)]
        )
        self.__y = self.crf.predict_single(self.__x_test)

        if pos:
            return [
                (self.__tokens[i], self.__pos_tags[i][1], data)
                for i, data in enumerate(self.__y)
            ]

        return [(self.__tokens[i], data) for i, data in enumerate(self.__y)]

    def __extract_features(self, doc):
        return [_doc2features(doc, i) for i in range(len(doc))]

    def __get_labels(self, doc):
        return [tag for (token, postag, tag) in doc]

    def __get_model(self):
        return self.crf
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Executable module to test unittest-xml-reporting.
"""
import sys

from xmlrunner.unittest import unittest
import xmlrunner
from six import StringIO, BytesIO
from tempfile import mkdtemp
from shutil import rmtree
from glob import glob
from xml.dom import minidom
import os.path

class XMLTestRunnerTestCase(unittest.TestCase):
    """
    XMLTestRunner test case.
    """
    class DummyTest(unittest.TestCase):
        @unittest.skip("demonstrating skipping")
        def test_skip(self):
            pass   # pragma: no cover
        @unittest.skip(u"demonstrating non-ascii skipping: éçà")
        def test_non_ascii_skip(self):
            pass   # pragma: no cover
        def test_pass(self):
            pass
        def test_fail(self):
            self.assertTrue(False)
        @unittest.expectedFailure
        def test_expected_failure(self):
            self.assertTrue(False)
        @unittest.expectedFailure
        def test_unexpected_success(self):
            pass
        def test_error(self):
            1 / 0
        def test_cdata_section(self):
            print('<![CDATA[content]]>')
        def test_non_ascii_error(self):
            self.assertEqual(u"éçà", 42)
        def test_unsafe_unicode(self):
            print(u"A\x00B\x08C\x0BD\x0C")
        def test_runner_buffer_output_pass(self):
            print('should not be printed')
        def test_runner_buffer_output_fail(self):
            print('should be printed')
            self.fail('expected to fail')
        def test_non_ascii_runner_buffer_output_fail(self):
            print(u'Where is the café ?')
            self.fail(u'The café could not be found')

    class DummySubTest(unittest.TestCase):
        def test_subTest_pass(self):
            for i in range(2):
                with self.subTest(i=i):
                    pass

        def test_subTest_fail(self):
            for i in range(2):
                with self.subTest(i=i):
                    self.fail('this is a subtest.')

    def setUp(self):
        self.stream = StringIO()
        self.outdir = mkdtemp()
        self.verbosity = 0
        self.runner_kwargs = {}
        self.addCleanup(rmtree, self.outdir)

    def _test_xmlrunner(self, suite, runner=None):
        outdir = self.outdir
        stream = self.stream
        verbosity = self.verbosity
        runner_kwargs = self.runner_kwargs
        if runner is None:
            runner = xmlrunner.XMLTestRunner(
                stream=stream, output=outdir, verbosity=verbosity,
                **self.runner_kwargs)
        self.assertEqual(0, len(glob(os.path.join(outdir, '*xml'))))
        runner.run(suite)
        self.assertEqual(1, len(glob(os.path.join(outdir, '*xml'))))
        return runner

    def test_basic_unittest_constructs(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        suite.addTest(self.DummyTest('test_skip'))
        suite.addTest(self.DummyTest('test_fail'))
        suite.addTest(self.DummyTest('test_expected_failure'))
        suite.addTest(self.DummyTest('test_unexpected_success'))
        suite.addTest(self.DummyTest('test_error'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_non_ascii(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_non_ascii_skip'))
        suite.addTest(self.DummyTest('test_non_ascii_error'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            u'<skipped message="demonstrating non-ascii skipping: éçà" type="skip"/>'.encode('utf8'),
            output)

    def test_xmlrunner_safe_xml_encoding_name(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        firstline = output.splitlines()[0]
        # test for issue #74
        self.assertIn('encoding="UTF-8"'.encode('utf8'), firstline)

    def test_xmlrunner_check_for_valid_xml_streamout(self):
        """
        This test checks if the xml document is valid if there are more than
        one testsuite and the output of the report is a single stream.
        """
        class DummyTestA(unittest.TestCase):
            def test_pass(self):
                pass
        class DummyTestB(unittest.TestCase):
            def test_pass(self):
                pass
        suite = unittest.TestSuite()
        suite.addTest( unittest.TestLoader().loadTestsFromTestCase(DummyTestA) );
        suite.addTest( unittest.TestLoader().loadTestsFromTestCase(DummyTestB) );
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        # Finally check if we have a valid XML document or not.
        try:
            minidom.parseString(output)
        except Exception as e:
            self.fail(e)

    def test_xmlrunner_unsafe_unicode(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_unsafe_unicode'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(u"<![CDATA[ABCD\n]]>".encode('utf8'), output)

    def test_xmlrunner_non_ascii_failures(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_non_ascii_runner_buffer_output_fail'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            u'Where is the café ?'.encode('utf8'),
            output)
        self.assertIn(
            u'The café could not be found'.encode('utf8'),
            output)

    @unittest.expectedFailure
    def test_xmlrunner_buffer_output_pass(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_runner_buffer_output_pass'))
        self._test_xmlrunner(suite)
        testsuite_output = self.stream.getvalue()
        # Since we are always buffering stdout/stderr
        # it is currently troublesome to print anything at all
        # and be consistent with --buffer option (issue #59)
        self.assertIn('should not be printed', testsuite_output)
        # this will be fixed when using the composite approach
        # that was under development in the rewrite branch.

    def test_xmlrunner_buffer_output_fail(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_runner_buffer_output_fail'))
        self._test_xmlrunner(suite)
        testsuite_output = self.stream.getvalue()
        self.assertIn('should be printed', testsuite_output)

    @unittest.skipIf(not hasattr(unittest.TestCase,'subTest'),
        'unittest.TestCase.subTest not present.')
    def test_unittest_subTest_fail(self):
        # test for issue #77
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummySubTest('test_subTest_fail'))
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            b'<testcase classname="tests.testsuite.DummySubTest" '
            b'name="test_subTest_fail (i=0)"',
            output)
        self.assertIn(
            b'<testcase classname="tests.testsuite.DummySubTest" '
            b'name="test_subTest_fail (i=1)"',
            output)

    @unittest.skipIf(not hasattr(unittest.TestCase, 'subTest'), 'unittest.TestCase.subTest not present.')
    def test_unittest_subTest_pass(self):
        # Test for issue #85
        suite = unittest.TestSuite()
        suite.addTest(self.DummySubTest('test_subTest_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_pass(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_verbose(self):
        self.verbosity = 1
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_showall(self):
        self.verbosity = 2
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_cdata_section(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_cdata_section'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_outsuffix(self):
        self.runner_kwargs['outsuffix'] = '.somesuffix'
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)
        xmlfile = glob(os.path.join(self.outdir, '*xml'))[0]
        assert xmlfile.endswith('.somesuffix.xml')

    def test_xmlrunner_nosuffix(self):
        self.runner_kwargs['outsuffix'] = ''
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)
        xmlfile = glob(os.path.join(self.outdir, '*xml'))[0]
        xmlfile = os.path.basename(xmlfile)
        assert xmlfile.endswith('DummyTest.xml')

    def test_junitxml_properties(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        suite.properties = dict(key='value')
        self._test_xmlrunner(suite)

    def test_xmlrunner_elapsed_times(self):
        self.runner_kwargs['elapsed_times'] = False
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_stream(self):
        stream = self.stream
        output = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=stream, output=output, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        runner.run(suite)

    def test_xmlrunner_output_subdir(self):
        stream = self.stream
        output = os.path.join(self.outdir, 'subdir')
        runner = xmlrunner.XMLTestRunner(
            stream=stream, output=output, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        runner.run(suite)

    def test_xmlrunner_patched_stdout(self):
        old_stdout, old_stderr = sys.stdout, sys.stderr
        try:
            sys.stdout, sys.stderr = StringIO(), StringIO()
            suite = unittest.TestSuite()
            suite.addTest(self.DummyTest('test_pass'))
            suite.properties = dict(key='value')
            self._test_xmlrunner(suite)
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr

<commit_msg>Add a test case for error added during __call__()<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Executable module to test unittest-xml-reporting.
"""
import sys

from xmlrunner.unittest import unittest
import xmlrunner
from six import StringIO, BytesIO
from tempfile import mkdtemp
from shutil import rmtree
from glob import glob
from xml.dom import minidom
import os.path

class XMLTestRunnerTestCase(unittest.TestCase):
    """
    XMLTestRunner test case.
    """
    class DummyTest(unittest.TestCase):
        @unittest.skip("demonstrating skipping")
        def test_skip(self):
            pass   # pragma: no cover
        @unittest.skip(u"demonstrating non-ascii skipping: éçà")
        def test_non_ascii_skip(self):
            pass   # pragma: no cover
        def test_pass(self):
            pass
        def test_fail(self):
            self.assertTrue(False)
        @unittest.expectedFailure
        def test_expected_failure(self):
            self.assertTrue(False)
        @unittest.expectedFailure
        def test_unexpected_success(self):
            pass
        def test_error(self):
            1 / 0
        def test_cdata_section(self):
            print('<![CDATA[content]]>')
        def test_non_ascii_error(self):
            self.assertEqual(u"éçà", 42)
        def test_unsafe_unicode(self):
            print(u"A\x00B\x08C\x0BD\x0C")
        def test_runner_buffer_output_pass(self):
            print('should not be printed')
        def test_runner_buffer_output_fail(self):
            print('should be printed')
            self.fail('expected to fail')
        def test_non_ascii_runner_buffer_output_fail(self):
            print(u'Where is the café ?')
            self.fail(u'The café could not be found')

    class DummySubTest(unittest.TestCase):
        def test_subTest_pass(self):
            for i in range(2):
                with self.subTest(i=i):
                    pass

        def test_subTest_fail(self):
            for i in range(2):
                with self.subTest(i=i):
                    self.fail('this is a subtest.')

    class DummyErrorInCallTest(unittest.TestCase):
        def __call__(self, result):
            try:
                raise Exception('Massive fail')
            except Exception:
                result.addError(self, sys.exc_info())
                return
            super(DummyErrorInCallTest, self).__call__(result)
        def test_pass(self):
            pass

    def setUp(self):
        self.stream = StringIO()
        self.outdir = mkdtemp()
        self.verbosity = 0
        self.runner_kwargs = {}
        self.addCleanup(rmtree, self.outdir)

    def _test_xmlrunner(self, suite, runner=None):
        outdir = self.outdir
        stream = self.stream
        verbosity = self.verbosity
        runner_kwargs = self.runner_kwargs
        if runner is None:
            runner = xmlrunner.XMLTestRunner(
                stream=stream, output=outdir, verbosity=verbosity,
                **self.runner_kwargs)
        self.assertEqual(0, len(glob(os.path.join(outdir, '*xml'))))
        runner.run(suite)
        self.assertEqual(1, len(glob(os.path.join(outdir, '*xml'))))
        return runner

    def test_basic_unittest_constructs(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        suite.addTest(self.DummyTest('test_skip'))
        suite.addTest(self.DummyTest('test_fail'))
        suite.addTest(self.DummyTest('test_expected_failure'))
        suite.addTest(self.DummyTest('test_unexpected_success'))
        suite.addTest(self.DummyTest('test_error'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_non_ascii(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_non_ascii_skip'))
        suite.addTest(self.DummyTest('test_non_ascii_error'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            u'<skipped message="demonstrating non-ascii skipping: éçà" type="skip"/>'.encode('utf8'),
            output)

    def test_xmlrunner_safe_xml_encoding_name(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        firstline = output.splitlines()[0]
        # test for issue #74
        self.assertIn('encoding="UTF-8"'.encode('utf8'), firstline)

    def test_xmlrunner_check_for_valid_xml_streamout(self):
        """
        This test checks if the xml document is valid if there are more than
        one testsuite and the output of the report is a single stream.
        """
        class DummyTestA(unittest.TestCase):
            def test_pass(self):
                pass
        class DummyTestB(unittest.TestCase):
            def test_pass(self):
                pass
        suite = unittest.TestSuite()
        suite.addTest( unittest.TestLoader().loadTestsFromTestCase(DummyTestA) );
        suite.addTest( unittest.TestLoader().loadTestsFromTestCase(DummyTestB) );
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        # Finally check if we have a valid XML document or not.
        try:
            minidom.parseString(output)
        except Exception as e:
            self.fail(e)

    def test_xmlrunner_unsafe_unicode(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_unsafe_unicode'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(u"<![CDATA[ABCD\n]]>".encode('utf8'), output)

    def test_xmlrunner_non_ascii_failures(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_non_ascii_runner_buffer_output_fail'))
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            u'Where is the café ?'.encode('utf8'),
            output)
        self.assertIn(
            u'The café could not be found'.encode('utf8'),
            output)

    @unittest.expectedFailure
    def test_xmlrunner_buffer_output_pass(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_runner_buffer_output_pass'))
        self._test_xmlrunner(suite)
        testsuite_output = self.stream.getvalue()
        # Since we are always buffering stdout/stderr
        # it is currently troublesome to print anything at all
        # and be consistent with --buffer option (issue #59)
        self.assertIn('should not be printed', testsuite_output)
        # this will be fixed when using the composite approach
        # that was under development in the rewrite branch.

    def test_xmlrunner_buffer_output_fail(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_runner_buffer_output_fail'))
        self._test_xmlrunner(suite)
        testsuite_output = self.stream.getvalue()
        self.assertIn('should be printed', testsuite_output)

    @unittest.skipIf(not hasattr(unittest.TestCase,'subTest'),
        'unittest.TestCase.subTest not present.')
    def test_unittest_subTest_fail(self):
        # test for issue #77
        outdir = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=self.stream, output=outdir, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummySubTest('test_subTest_fail'))
        runner.run(suite)
        outdir.seek(0)
        output = outdir.read()
        self.assertIn(
            b'<testcase classname="tests.testsuite.DummySubTest" '
            b'name="test_subTest_fail (i=0)"',
            output)
        self.assertIn(
            b'<testcase classname="tests.testsuite.DummySubTest" '
            b'name="test_subTest_fail (i=1)"',
            output)

    @unittest.skipIf(not hasattr(unittest.TestCase, 'subTest'), 'unittest.TestCase.subTest not present.')
    def test_unittest_subTest_pass(self):
        # Test for issue #85
        suite = unittest.TestSuite()
        suite.addTest(self.DummySubTest('test_subTest_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_pass(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_verbose(self):
        self.verbosity = 1
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_showall(self):
        self.verbosity = 2
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_cdata_section(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_cdata_section'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_outsuffix(self):
        self.runner_kwargs['outsuffix'] = '.somesuffix'
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)
        xmlfile = glob(os.path.join(self.outdir, '*xml'))[0]
        assert xmlfile.endswith('.somesuffix.xml')

    def test_xmlrunner_nosuffix(self):
        self.runner_kwargs['outsuffix'] = ''
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)
        xmlfile = glob(os.path.join(self.outdir, '*xml'))[0]
        xmlfile = os.path.basename(xmlfile)
        assert xmlfile.endswith('DummyTest.xml')

    def test_junitxml_properties(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        suite.properties = dict(key='value')
        self._test_xmlrunner(suite)

    def test_xmlrunner_elapsed_times(self):
        self.runner_kwargs['elapsed_times'] = False
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        self._test_xmlrunner(suite)

    def test_xmlrunner_stream(self):
        stream = self.stream
        output = BytesIO()
        runner = xmlrunner.XMLTestRunner(
            stream=stream, output=output, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        runner.run(suite)

    def test_xmlrunner_output_subdir(self):
        stream = self.stream
        output = os.path.join(self.outdir, 'subdir')
        runner = xmlrunner.XMLTestRunner(
            stream=stream, output=output, verbosity=self.verbosity,
            **self.runner_kwargs)
        suite = unittest.TestSuite()
        suite.addTest(self.DummyTest('test_pass'))
        runner.run(suite)

    def test_xmlrunner_patched_stdout(self):
        old_stdout, old_stderr = sys.stdout, sys.stderr
        try:
            sys.stdout, sys.stderr = StringIO(), StringIO()
            suite = unittest.TestSuite()
            suite.addTest(self.DummyTest('test_pass'))
            suite.properties = dict(key='value')
            self._test_xmlrunner(suite)
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr

    def test_xmlrunner_error_in_call(self):
        suite = unittest.TestSuite()
        suite.addTest(self.DummyErrorInCallTest('test_pass'))
        self._test_xmlrunner(suite)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import sys
sys.path.append("..")
import jauxiliar as jaux
import josecommon as jcommon
import discord

MARKDOWN_HELP_FILES = [
    'doc/cmd/listcmd.md',
    'doc/cmd/admin.md',
    'doc/josespeak.md',
    'doc/jcoin.md'
]

JOSE_INVITE = 'https://discord.gg/5ASwg4C'

class JoseHelp(jaux.Auxiliar):
    def __init__(self, _client):
        jaux.Auxiliar.__init__(self, _client)
        self.help = {}

    def load_helpfiles(self):
        for file in MARKDOWN_HELP_FILES:
            with open(file, 'r') as f:
                # do I really need to parse MARKDOWN, without libs?
                for line in f.readlines():
                    if line.find('|'):
                        data = [sp.strip() for sp in line.split('|')]
                        if len(data) == 4:
                            raw_cmd = data[0]
                            c = raw_cmd.find('`')
                            a_cmd = raw_cmd[c+1:raw_cmd.find('`', c+1)]
                            a_cmd = a_cmd.split()
                            cmd = '`%s`' % (a_cmd[0])
                            desc = data[1]
                            examples = data[2]
                            aliases = data[3]

                            self.help[cmd] = {
                                'description': desc,
                                'examples': examples,
                                'aliases': aliases
                            }

    async def ext_load(self):
        self.load_helpfiles()
        return True, ''

    async def ext_unload(self):
        del self.help
        return True, ''

    async def c_helpme(self, message, args, cxt):
        '''`j!helpme` - having problems?'''

        if message.server.id != jcommon.JOSE_DEV_SERVER_ID:
            em = discord.Embed(title="Having Problems?", colour=discord.Colour.red())
            em.add_field("José server: ", "{}".format(JOSE_INVITE))
            em.add_field("what do", "Use `j!helpme` there and find the nearest mod that works with your issue")
            await cxt.say_embed(em)
            return

        res = []
        res.append(":large_blue_circle: online admins, :white_circle: idle admins")

        for adminid in jcommon.ADMIN_TOPICS:
            topics = jcommon.ADMIN_TOPICS[adminid]
            admin = message.server.get_member(adminid)
            if admin.status == discord.Status.online:
                res.append(":large_blue_circle: %s, works on %s" % (admin, ', '.join(topics)))
            elif admin.status == discord.Status.idle:
                res.append(":white_circle: %s, works on %s" % (admin, ', '.join(topics)))
            else:
                # admin is Do Not Disturb or Offline, don't show them lol
                pass

        await cxt.say('\n'.join(res))

    async def c_help(self, message, args, cxt):
        '''`j!help cmd` - get help for a command, uses the Markdown helptexts'''
        if len(args) < 2:
            await cxt.say(self.c_help.__doc__)
            return

        command = args[1]

        # thats ugly, just... ugly.
        command = '`j!%s`' % command

        if command in self.help:
            res = []
            helpdata = self.help[command]

            res.append("Command: `%s`" % command)
            res.append("**Description**: %s" % helpdata['description'])
            res.append("Example: %s" % helpdata['examples'])
            res.append("Aliases: %s" % helpdata['aliases'])

            await cxt.say('\n'.join(res))
        else:
            await cxt.say("No helptext was found for `%s`, try `j!docstring`", (command,))
        return

    async def c_docstring(self, message, args, cxt):
        '''`j!docstring command` - get the docstring for that command'''

        # load helptext
        cmd_ht = 'docstring'
        try:
            if args[1] == 'docstring':
                await cxt.say(self.c_docstring.__doc__)
                return
            else:
                cmd_ht = args[1]
        except:
            pass

        if cmd_ht == 'docstring':
            await cxt.say(self.c_docstring.__doc__)
            return

        cmd_method = getattr(cxt.jose, 'c_%s' % cmd_ht, None)
        if cmd_method is None:
            await cxt.say("%s: Command not found" % cmd_ht)
            return

        try:
            docstring = cmd_method.__doc__
            if docstring is None:
                await cxt.say("Command found, docstring not found")
            else:
                await cxt.say(docstring)
        except Exception as e:
            await cxt.say("Error getting docstring for %s: %r" % (cmd_ht, repr(e)))
<commit_msg>jhelp: helpme: use name and value in em.add_field<commit_after>#!/usr/bin/env python3

import sys
sys.path.append("..")
import jauxiliar as jaux
import josecommon as jcommon
import discord

MARKDOWN_HELP_FILES = [
    'doc/cmd/listcmd.md',
    'doc/cmd/admin.md',
    'doc/josespeak.md',
    'doc/jcoin.md'
]

JOSE_INVITE = 'https://discord.gg/5ASwg4C'

class JoseHelp(jaux.Auxiliar):
    def __init__(self, _client):
        jaux.Auxiliar.__init__(self, _client)
        self.help = {}

    def load_helpfiles(self):
        for file in MARKDOWN_HELP_FILES:
            with open(file, 'r') as f:
                # do I really need to parse MARKDOWN, without libs?
                for line in f.readlines():
                    if line.find('|'):
                        data = [sp.strip() for sp in line.split('|')]
                        if len(data) == 4:
                            raw_cmd = data[0]
                            c = raw_cmd.find('`')
                            a_cmd = raw_cmd[c+1:raw_cmd.find('`', c+1)]
                            a_cmd = a_cmd.split()
                            cmd = '`%s`' % (a_cmd[0])
                            desc = data[1]
                            examples = data[2]
                            aliases = data[3]

                            self.help[cmd] = {
                                'description': desc,
                                'examples': examples,
                                'aliases': aliases
                            }

    async def ext_load(self):
        self.load_helpfiles()
        return True, ''

    async def ext_unload(self):
        del self.help
        return True, ''

    async def c_helpme(self, message, args, cxt):
        '''`j!helpme` - having problems?'''

        if message.server.id != jcommon.JOSE_DEV_SERVER_ID:
            em = discord.Embed(title="Having Problems?", colour=discord.Colour.red())
            em.add_field(name="José server: ", value="{}".format(JOSE_INVITE))
            em.add_field(name="what do", value="Use `j!helpme` there and find the nearest mod that works with your issue")
            await cxt.say_embed(em)
            return

        res = []
        res.append(":large_blue_circle: online admins, :white_circle: idle admins")

        for adminid in jcommon.ADMIN_TOPICS:
            topics = jcommon.ADMIN_TOPICS[adminid]
            admin = message.server.get_member(adminid)
            if admin.status == discord.Status.online:
                res.append(":large_blue_circle: %s, works on %s" % (admin, ', '.join(topics)))
            elif admin.status == discord.Status.idle:
                res.append(":white_circle: %s, works on %s" % (admin, ', '.join(topics)))
            else:
                # admin is Do Not Disturb or Offline, don't show them lol
                pass

        await cxt.say('\n'.join(res))

    async def c_help(self, message, args, cxt):
        '''`j!help cmd` - get help for a command, uses the Markdown helptexts'''
        if len(args) < 2:
            await cxt.say(self.c_help.__doc__)
            return

        command = args[1]

        # thats ugly, just... ugly.
        command = '`j!%s`' % command

        if command in self.help:
            res = []
            helpdata = self.help[command]

            res.append("Command: `%s`" % command)
            res.append("**Description**: %s" % helpdata['description'])
            res.append("Example: %s" % helpdata['examples'])
            res.append("Aliases: %s" % helpdata['aliases'])

            await cxt.say('\n'.join(res))
        else:
            await cxt.say("No helptext was found for `%s`, try `j!docstring`", (command,))
        return

    async def c_docstring(self, message, args, cxt):
        '''`j!docstring command` - get the docstring for that command'''

        # load helptext
        cmd_ht = 'docstring'
        try:
            if args[1] == 'docstring':
                await cxt.say(self.c_docstring.__doc__)
                return
            else:
                cmd_ht = args[1]
        except:
            pass

        if cmd_ht == 'docstring':
            await cxt.say(self.c_docstring.__doc__)
            return

        cmd_method = getattr(cxt.jose, 'c_%s' % cmd_ht, None)
        if cmd_method is None:
            await cxt.say("%s: Command not found" % cmd_ht)
            return

        try:
            docstring = cmd_method.__doc__
            if docstring is None:
                await cxt.say("Command found, docstring not found")
            else:
                await cxt.say(docstring)
        except Exception as e:
            await cxt.say("Error getting docstring for %s: %r" % (cmd_ht, repr(e)))
<|endoftext|>


<commit_before>#!/usr/bin/env python

#    /files/some/path/stats/json

from __future__ import print_function
import requests
import json
import datetime

import pandas as pd
import numpy as np

import cartopy.crs as ccrs
import cartopy.feature as cf
import cartopy.io.shapereader as shpreader

import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.patches as mpatches

mpl.style.use('ggplot')

def get_sourceforge_stats(project,
                          start_date,
                          end_date=datetime.date.today(),
                          path="/"):

   """
   """

   # Get download statistics from SourceForge
   #
   # The API uses URLs very similar to the existing stats URLs. Here's an example:
   #    https://sourceforge.net/projects/openofficeorg.mirror/files/stats/json?start_date=2014-10-29&end_date=2014-11-04
   #
   # The url parameters start_date and end_date are required and must be
   # dates in YYYY-MM-DD format. To access project-wide stats, use
   #    /files/stats/json in the URL. To limit to a file or folder, use
   #    /files/some/path/stats/json

   print("start_date: {0}".format(start_date))
   query_url="https://sourceforge.net/projects/{0}/files/{1}/stats/json?start_date={2}&end_date={3}".format(project,
                                                                                path,
                                                                                start_date.isoformat(),
                                                                                end_date.isoformat())
   print("Query: \"{0}\"...".format(query_url))
   r = requests.get(query_url)

   if r.status_code == requests.codes.ok:
       return json.loads(r.text)
   else:
       r.raise_for_status()

#project="itk"
#path="itk"

project="simpleitk"
path="SimpleITK"
start_date=datetime.date(2013,1,1)
end_date=datetime.date.today() - datetime.timedelta(days=1)

# Load the JSON data into a nested dictionary structure
sf_json = get_sourceforge_stats(project,
                                start_date,
                                end_date,
                                path)

# Load the downloads over time into a Pandas Series
sf_dbym = pd.DataFrame(sf_json["downloads"])
sf_dbym.columns=["Date","Download Count"]
sf_dbym.set_index(sf_dbym.columns[0],inplace=True)
sf_dbym.index = pd.to_datetime(sf_dbym.index)

sf_dbym.plot(kind='area')
plt.title("SimpleITK SourceForge Downloads")
plt.ylabel("Downloads per Month")
plt.show()

# load the country, OS download counts into a DataFrame, from json
# the json field "oses_by_country" is a dict of countries indexing a dict of OSes and download
cbo_df = pd.DataFrame([v for (k,v) in sf_json["oses_by_country"].items()], index=sf_json["oses_by_country"].keys())
cbo_df.index.name="Country"
cbo_df.fillna(0,inplace=True)
cbo_df = cbo_df.astype(int)

# sort the data frame with the most popular country and OS at 0 (upper left)
cbo_df = cbo_df.loc[cbo_df.sum(axis=1).sort_values(ascending=False).index]
cbo_df = cbo_df[cbo_df.sum(axis=0).sort_values(ascending=False).index]
print(cbo_df)

# Graph OSes
# TODO: Combine OSes into a "Other" group
os_series= cbo_df.sum(axis=0).sort_values(ascending=False)
os_series.name = "OS"
os_series.plot(kind='pie',
               autopct='%1.1f%%',
               title="SourceForge Downloads by OS")
plt.show()

# Graph countries with pie Chart
total = cbo_df.sum(axis=1).sum()
percent = 0.80
bool_index = cbo_df.sum(axis=1).sort_values(ascending=False).cumsum()/total<percent
df = cbo_df.loc[bool_index].copy()
df.loc["Other"] = cbo_df.loc[~(bool_index)].sum(axis=0)
print(df)
df.sum(axis=1).plot(kind='pie',
            autopct='%1.1f%%',
            title="Downloads by Country")
plt.show();

#
# Cartography plot of downloads by country
#
c_series = cbo_df.sum(axis=1)
c_series = c_series.sort_values(ascending=False)
cmap = mpl.cm.OrRd

max_dl = float(c_series.max())
min_dl = float(c_series.min())

norm = mpl.colors.LogNorm(vmin=min_dl, vmax=max_dl)
total_users = c_series.sum()
shapename = 'admin_0_countries'
countries_shp = shpreader.natural_earth(resolution='110m', category='cultural', name=shapename)

plt.figure(figsize=(14,8))
ax = plt.axes(projection=ccrs.Robinson() )
ax.add_feature(cf.BORDERS)
ax.add_feature(cf.COASTLINE)
#ax.add_feature(cf.OCEAN)
#ax.add_feature(cf.LAND)

hdl_dict = dict()

for country in shpreader.Reader(countries_shp).records():
    name = country.attributes['name_long']
    if name in c_series:
        num_users = c_series[name]
        hdl = ax.add_geometries(country.geometry, ccrs.PlateCarree(),
                               facecolor=cmap(norm(num_users),alpha=True))
        hdl_dict[name] = hdl

# Add ColorBar
sm = plt.cm.ScalarMappable(cmap=cmap,norm=norm)
sm._A = [] # WHAT? DANGER
plt.colorbar(sm,ax=ax)
plt.title("Source Forge Downloads by Country\n{0} to {1}".format(start_date, end_date))

# Add Legend
cnt = (cbo_df.sum(axis=1).sort_values(ascending=False).cumsum()/total < percent).sum()

p = [ mpatches.Patch(color=cmap(norm(v)), label="{0} ({1})".format(k,v)) for (k,v) in c_series[0:cnt].iteritems()]
plt.legend(handles=p,
           loc=3, fancybox=True)
plt.savefig("{0}_CountriesCarto_{1}_{2}.pdf".format(project,start_date, end_date), format='pdf')
plt.show()

##################################################
#
# PyPI Download Statistics - via Google BigData Query
#
#
"""
SELECT
  STRFTIME_UTC_USEC(timestamp, "%Y-%m") AS yyyymm,
  COUNT(*) as download_count
FROM
  TABLE_DATE_RANGE(
    [the-psf:pypi.downloads],
    DATE_ADD(CURRENT_TIMESTAMP(), -12, "month"),
    CURRENT_TIMESTAMP()
  )
WHERE
  file.project = 'simpleitk' and details.installer.name != 'bandersnatch'
GROUP BY
  yyyymm, file.version,
ORDER BY
  yyyymm DESC, download_count DESC
LIMIT 100
"""
#
##################################################

pypi_ = pd.read_csv("psf-pypi-downloads-simpleitk-query.csv")
print(df)

# Load query results into a data frame
df["yyyymm"] = pd.to_datetime(df["yyyymm"], format=("%Y-%m") )
df = df.pivot(index='yyyymm', columns='file_version', values='download_count')
df.index.name = "Date"
df.columns.name = "Version"
df.fillna(0,inplace=True)
df = df.astype(int)
print(df)
print("Total PyPI Downloads: {0}".format(df.sum(axis=1).sum()))

# Plot the PyPI downloads of different versions
df.plot(kind="area",
           title="SimpleITK PyPI Downloads")
plt.ylabel("Downloads per Month")
plt.show()

# Create Series for PyPI downloads by time
s = df.sum(axis=1)
s.index.name="Date"
s.name= "PyPI Download Count"

# Manually enter dates of tags
tags = pd.Series({"2017-04-05":"v1.0.0",
                "2016-07-27":"v0.10.0",
                  "2015-09-25":"v0.9.1",
                  "2015-05-012":"v0.9.0",
                  "2014-12-19":"v0.8.1",
                  "2014-03-8":"v0.8.0",
                  "2013-11-21":"v0.7.1",
                  "2013-09-30":"v0.7.0",
                  "2013-02-1":"v0.6.0"})
tags.index = pd.to_datetime(tags.index)
tags.name = "Releases"

#
# Combine SourceForge and PyPI download statistics into Total Downloads
#
combined_df = pd.concat([sf_dbym,pd.DataFrame(s)],axis=1)
combined_df.columns = ["Source Forge Downloads", combined_df.columns[1]]

plt.figure(figsize=(8,10))
ax = combined_df.plot(kind="area")

i = 1
for (date,tag) in tags.iteritems():
    #ax.plot(date, 0, 'o')
    ax.annotate(tag,(date,800), horizontalalignment='left', verticalalignment='center', rotation=90)
    i = i +1

plt.title("SimpleITK Downloads");
plt.ylabel("Downloads per Month")

plt.savefig("{0}_Downloads_{1}_{2}.pdf".format(project,start_date, end_date), format='pdf')
plt.show()
<commit_msg>Update download statistics script<commit_after>#!/usr/bin/env python

#    /files/some/path/stats/json

from __future__ import print_function
import requests
import json
import datetime

import pandas as pd
import numpy as np

import cartopy.crs as ccrs
import cartopy.feature as cf
import cartopy.io.shapereader as shpreader

import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.patches as mpatches

mpl.style.use('ggplot')

def get_sourceforge_stats(project,
                          start_date,
                          end_date=datetime.date.today(),
                          path="/"):

   """
   """

   # Get download statistics from SourceForge
   #
   # The API uses URLs very similar to the existing stats URLs. Here's an example:
   #    https://sourceforge.net/projects/openofficeorg.mirror/files/stats/json?start_date=2014-10-29&end_date=2014-11-04
   #
   # The url parameters start_date and end_date are required and must be
   # dates in YYYY-MM-DD format. To access project-wide stats, use
   #    /files/stats/json in the URL. To limit to a file or folder, use
   #    /files/some/path/stats/json

   print("start_date: {0}".format(start_date))
   query_url="https://sourceforge.net/projects/{0}/files/{1}/stats/json?start_date={2}&end_date={3}".format(project,
                                                                                path,
                                                                                start_date.isoformat(),
                                                                                end_date.isoformat())
   print("Query: \"{0}\"...".format(query_url))
   r = requests.get(query_url)

   if r.status_code == requests.codes.ok:
       return json.loads(r.text)
   else:
       r.raise_for_status()

#project="itk"
#path="itk"

project="simpleitk"
path="SimpleITK"
start_date=datetime.date(2013,1,1)
end_date=datetime.date.today() - datetime.timedelta(days=1)

# Load the JSON data into a nested dictionary structure
sf_json = get_sourceforge_stats(project,
                                start_date,
                                end_date,
                                path)

# Load the downloads over time into a Pandas Series
sf_dbym = pd.DataFrame(sf_json["downloads"])
sf_dbym.columns=["Date","Download Count"]
sf_dbym.set_index(sf_dbym.columns[0],inplace=True)
sf_dbym.index = pd.to_datetime(sf_dbym.index)

sf_dbym.plot(kind='area')
plt.title("SimpleITK SourceForge Downloads")
plt.ylabel("Downloads per Month")
plt.show()

# load the country, OS download counts into a DataFrame, from json
# the json field "oses_by_country" is a dict of countries indexing a dict of OSes and download
cbo_df = pd.DataFrame([v for (k,v) in sf_json["oses_by_country"].items()], index=sf_json["oses_by_country"].keys())
cbo_df.index.name="Country"
cbo_df.fillna(0,inplace=True)
cbo_df = cbo_df.astype(int)

# sort the data frame with the most popular country and OS at 0 (upper left)
cbo_df = cbo_df.loc[cbo_df.sum(axis=1).sort_values(ascending=False).index]
cbo_df = cbo_df[cbo_df.sum(axis=0).sort_values(ascending=False).index]
print(cbo_df)

# Graph OSes
# TODO: Combine OSes into a "Other" group
os_series= cbo_df.sum(axis=0).sort_values(ascending=False)
os_series.name = "OS"
os_series.plot(kind='pie',
               autopct='%1.1f%%',
               title="SourceForge Downloads by OS")
plt.show()

# Graph countries with pie Chart
total = cbo_df.sum(axis=1).sum()
percent = 0.80
bool_index = cbo_df.sum(axis=1).sort_values(ascending=False).cumsum()/total<percent
df = cbo_df.loc[bool_index].copy()
df.loc["Other"] = cbo_df.loc[~(bool_index)].sum(axis=0)
print(df)
df.sum(axis=1).plot(kind='pie',
            autopct='%1.1f%%',
            title="Downloads by Country")
plt.show();

#
# Cartography plot of downloads by country
#
c_series = cbo_df.sum(axis=1)
c_series = c_series.sort_values(ascending=False)
cmap = mpl.cm.OrRd

max_dl = float(c_series.max())
min_dl = float(c_series.min())

norm = mpl.colors.LogNorm(vmin=min_dl, vmax=max_dl)
total_users = c_series.sum()
shapename = 'admin_0_countries'
countries_shp = shpreader.natural_earth(resolution='110m', category='cultural', name=shapename)

plt.figure(figsize=(14,8))
ax = plt.axes(projection=ccrs.Robinson() )
ax.add_feature(cf.BORDERS)
ax.add_feature(cf.COASTLINE)
#ax.add_feature(cf.OCEAN)
#ax.add_feature(cf.LAND)

hdl_dict = dict()

for country in shpreader.Reader(countries_shp).records():
    name = country.attributes['name_long']
    if name in c_series:
        num_users = c_series[name]
        hdl = ax.add_geometries(country.geometry, ccrs.PlateCarree(),
                               facecolor=cmap(norm(num_users),alpha=True))
        hdl_dict[name] = hdl

# Add ColorBar
sm = plt.cm.ScalarMappable(cmap=cmap,norm=norm)
sm._A = [] # WHAT? DANGER
plt.colorbar(sm,ax=ax)
plt.title("Source Forge Downloads by Country\n{0} to {1}".format(start_date, end_date))

# Add Legend
cnt = (cbo_df.sum(axis=1).sort_values(ascending=False).cumsum()/total < percent).sum()

p = [ mpatches.Patch(color=cmap(norm(v)), label="{0} ({1})".format(k,v)) for (k,v) in c_series[0:cnt].iteritems()]
plt.legend(handles=p,
           loc=3, fancybox=True)
plt.savefig("{0}_CountriesCarto_{1}_{2}.pdf".format(project,start_date, end_date), format='pdf')
plt.show()

##################################################
#
# PyPI Download Statistics - via Google BigData Query
#
#
"""
SELECT
  STRFTIME_UTC_USEC(timestamp, "%Y-%m") AS yyyymm,
  COUNT(*) as download_count
FROM
  TABLE_DATE_RANGE(
    [the-psf:pypi.downloads],
    DATE_ADD(CURRENT_TIMESTAMP(), -12, "month"),
    CURRENT_TIMESTAMP()
  )
WHERE
  file.project = 'simpleitk' and details.installer.name != 'bandersnatch'
GROUP BY
  yyyymm, file.version
ORDER BY
  yyyymm DESC, download_count DESC, file.version DESC
LIMIT 100
"""
#
##################################################

df = pd.read_csv("psf-pypi-downloads-simpleitk-query.csv")
print(df)

# Load query results into a data frame
df["yyyymm"] = pd.to_datetime(df["yyyymm"], format=("%Y-%m") )
if "file_version" in df.columns:
   df = df.pivot(index='yyyymm', columns='file_version', values='download_count')
   df.columns.name = "Version"
else:
   df.set_index('yyyymm', inplace=True)

df.fillna(0,inplace=True)
df = df.astype(int)
df.sort_index(ascending=False,inplace=True)

print(df)
print("Total PyPI Downloads: {0}".format(df.sum(axis=1).sum()))

# Plot the PyPI downloads of different versions
df.plot(kind="area",
           title="SimpleITK PyPI Downloads")
plt.ylabel("Downloads per Month")
plt.show()

# Create Series for PyPI downloads by time
s = df.sum(axis=1)
s.name= "PyPI Download Count"
s.index.name="Date"

# Manually enter dates of tags
tags = pd.Series({"2017-04-05":"v1.0.0",
                  "2016-07-27":"v0.10.0",
                  "2015-09-25":"v0.9.1",
                  "2015-05-12":"v0.9.0",
                  "2014-12-19":"v0.8.1",
                  "2014-03-08":"v0.8.0",
                  "2013-11-21":"v0.7.1",
                  "2013-09-30":"v0.7.0",
                  "2013-02-01":"v0.6.0"})
tags.index = pd.to_datetime(tags.index)
tags.name = "Releases"

#
# Combine SourceForge and PyPI download statistics into Total Downloads
#
combined_df = pd.concat([sf_dbym,pd.DataFrame(s)],axis=1)
combined_df.columns = ["Source Forge Downloads", combined_df.columns[1]]

plt.figure(figsize=(8,10))
ax = combined_df.plot(kind="area")

i = 1
for (date,tag) in tags.iteritems():
    #ax.plot(date, 0, 'o')
    ax.annotate(tag,(date,800), horizontalalignment='left', verticalalignment='center', rotation=90)
    i = i +1

plt.title("SimpleITK Downloads");
plt.ylabel("Downloads per Month")

plt.savefig("{0}_Downloads_{1}_{2}.pdf".format(project,start_date, end_date), format='pdf')
plt.show()
<|endoftext|>