<commit_before>from PyQt5 import QtGui
from region_selector_ui import Ui_RegionSelector
import pyqtgraph as pg
import numpy as np
from PIL import Image

from utilities import RegionROI

class RegionSelector(QtGui.QMainWindow):
    """
    User interface logic to select stack regions
    """

    def __init__(self, parent=None):
        """
        Creates a new RegionSelector window
        :param parent: The window's parent
        """
        # init of ui
        QtGui.QWidget.__init__(self, parent)
        self.ui = Ui_RegionSelector()
        self.ui.setupUi(self)
        # non-ui class members
        self.__current_z = 0
        self.filename = ""
        self.currentStack = np.array([])
        self.roi_dict = {}
        self.current_z = 0
        self.current_ROI = None
        self.last_color = 0
        self.last_uid = 0
        # ui stuff settings
        self.ui.lblROIName.setText("")
        # create our view-box and image view inside stackBox
        self.stack_vbox = self.ui.stackBox.addViewBox()
        self.stack_vbox.invertY(True)  # since image display inverts the Y-axis
        self.stack_vbox.setAspectLocked(True)
        self.stack_image = pg.ImageItem()
        self.stack_vbox.addItem(self.stack_image)
        # hide menu and roi button on image views
        self.ui.ROIView.ui.roiBtn.hide()
        self.ui.ROIView.ui.menuBtn.hide()
        # connect signals
        self.ui.btnLoad.clicked.connect(self.load_click)
        self.ui.btnAddROI.clicked.connect(self.addroi_click)
        self.ui.cbRegions.currentIndexChanged.connect(self.regionNameSelChanged)
        self.ui.sldrZ.sliderMoved.connect(self.sliderZChanged)

        # for easy testing, create and display test-image stack
        test_image = np.abs(np.random.randn(10, 100, 100, 3)*50).astype(np.uint8)
        test_image[:, 50, :, :] = 255
        test_image[:, :, 50, :] = 255
        self.currentStack = test_image
        self.reset_after_load()
        self.display_slice()

    def guess_stack_type(self):
        """
        Uses simple heuristics to determine the type (2D, 3D, grayscale, color) of the current stack
        :return: Integer identifying the stack type
        """
        if self.currentStack.size == 0 or self.currentStack.ndim < 2:
            # stack invalid
            return -1
        if self.currentStack.ndim == 2:
            # grayscale 2D image
            return 0
        if self.currentStack.ndim == 3:
            if self.currentStack.shape[2] <= 4:
                # RGB(A) 2D image (or very thin stack...)
                return 1
            else:
                # greyscale 3D stack
                return 2
        if self.currentStack.ndim == 4:
            if self.currentStack.shape[3] > 4:
                # invalid color model
                return -1
            else:
                # color 3D stack
                return 3
        # no idea
        return -1

    def reset_after_load(self):
        """
        Resets user-interface elements and internal counters after a new stack has been loaded
        """
        self.roi_dict = {}
        self.current_z = 0
        self.current_ROI = None
        self.last_color = 0
        self.last_uid = 0
        self.ui.cbRegions.clear()
        self.ui.sldrZ.setMinimum(0)
        self.ui.sldrZ.setMaximum(self.NSlices - 1)
        self.ui.sldrZ.setValue(0)

    @property
    def NSlices(self):
        """
        The number of slices in the currently loaded stack
        """
        st = self.guess_stack_type()
        if st == -1:
            return 0
        if st < 2:
            return 1
        return self.currentStack.shape[0]

    @property
    def current_z(self):
        return self.__current_z

    @current_z.setter
    def current_z(self, current_z):
        if current_z >= 0:
            self.__current_z = current_z
            self.ui.lbl_z.setText(str(current_z))

    def display_slice(self):
        """
        Displays the current slice
        """
        st = self.guess_stack_type()
        if st == -1:
            print("No valid stack loaded")
            return
        if st == 0 or st == 1:
            self.stack_image.setImage(self.currentStack)
        elif st == 2 or st == 3:
            if self.NSlices > self.current_z:
                self.stack_image.setImage(self.currentStack[self.current_z])
            else:
                print("Improper slice index for stack")

    def create_default_roi(self, n_vert=6):
        """
        Create a point list for the default roi when adding a new roi
        :param n_vert: The initial number of vertices in the ROI
        :return: A list of (x,y) tuples for the vertices of the roi
        """
        st = self.guess_stack_type()
        if st == -1:
            return
        if st < 2:
            width = self.currentStack.shape[0]
            height = self.currentStack.shape[1]
        else:
            width = self.currentStack.shape[1]
            height = self.currentStack.shape[2]
        center_x = width / 2
        center_y = height / 2
        radius = min([width, height]) / 10
        point_list = []
        for i in range(n_vert):
            a = 2*np.pi / n_vert * i
            x = np.cos(a)*radius + center_x
            y = np.sin(a)*radius + center_y
            point_list.append((x, y))
        return point_list

    def next_roi_color(self):
        """
        Updates our ROI pen color cycle
        :return: The color of the next roi
        """
        self.last_color += 1
        self.last_color %= 9
        return self.last_color

    def next_roi_uid(self):
        """
        Updates the unique id of roi's
        :return: The uid of the next roi
        """
        self.last_uid += 1
        return self.last_uid

    # Signals #
    def load_click(self):
        """
        Handles event of clicking the load stack button
        """
        diag = QtGui.QFileDialog()
        fname = diag.getOpenFileName(self, "Select stack", "E:/Dropbox/2P_Data", "*.tif")[0]
        if fname is not None and fname != "":
            assert isinstance(fname, str)
            self.currentStack = self.OpenStack(fname)
            self.filename = fname
            self.reset_after_load()
            self.display_slice()
        else:
            print("No file selected")

    def addroi_click(self):
        """
        Handles event of clicking the add ROI button
        """
        name = self.ui.leNewROI.text()
        new_r = RegionROI(self.create_default_roi(), self.next_roi_uid(), name, pen=(self.next_roi_color(), 12))
        if self.current_z in self.roi_dict:
            self.roi_dict[self.current_z].append(new_r)
        else:
            self.roi_dict[self.current_z] = [new_r]
        # add the name of the created region to our list if it doesn exist yet
        found = False
        for i in range(self.ui.cbRegions.count()):
            if self.ui.cbRegions.itemText(i) == new_r.region_name:
                found = True
        if not found:
            self.ui.cbRegions.addItem(new_r.region_name)
        # add the ROI to the display
        self.stack_vbox.addItem(new_r)
        self.updateRoi(new_r)
        new_r.sigRegionChanged.connect(self.updateRoi)
        # clear the name field and combo box selection
        self.ui.leNewROI.setText("")
        self.ui.cbRegions.setCurrentIndex(-1)

    def updateRoi(self, roi: RegionROI):
        """
        Handles the event whenever an ROI is updated on-screen or programmatically
        :param roi: The roi which got updated
        """
        if roi is None:
            # clear our details image and roi indicator
            self.ui.ROIView.getImageItem().setImage(np.zeros((10,10)))
            self.ui.lblROIName.setText("")
            return
        self.current_ROI = roi
        self.ui.lblROIName.setText(roi.region_name)
        arr = roi.getArrayRegion(self.stack_image.image, img=self.stack_image)
        self.ui.ROIView.getImageItem().setImage(arr)

    def regionNameSelChanged(self, index):
        """
        Handles the event whenever a new region is selected in the combo box of previous region names
        :param index: The new index selection
        """
        if index >= 0:
            self.ui.leNewROI.setText(self.ui.cbRegions.itemText(index))

    def sliderZChanged(self, value):
        """
        Handles the event when our z-position slider was adjusted
        :param value: The new slider position
        """
        # decommission al ROI's of the current plane
        if self.current_z in self.roi_dict:
            for r in self.roi_dict[self.current_z]:
                self.stack_vbox.removeItem(r)
        # update z-plane
        self.current_z = value
        # load any existing ROI's of that plane and select the first
        if self.current_z in self.roi_dict:
            for i, r in enumerate(self.roi_dict[self.current_z]):
                self.stack_vbox.addItem(r)
                if i == 0:
                    self.updateRoi(r)
        else:
            self.updateRoi(None)
        # display the new slice
        self.display_slice()

    @staticmethod
    def OpenStack(filename):
        """
        Load image stack from tiff file
        :param filename: The name of the tiff file
        :return: A (z, x, y) numpy array
        """
        im = Image.open(filename)
        # NOTE: PIL transposes x and y coordinates on image loading, i.e. x coordinates
        # are increasing with increasing rows not columns
        if im.mode == 'RGB':
            stack = np.empty((im.n_frames, im.size[1], im.size[0], 3), dtype=np.uint8)
        else:
            stack = np.empty((im.n_frames, im.size[1], im.size[0]), dtype=np.uint8)
        # loop over frames and assign
        for i in range(im.n_frames):
            im.seek(i)
            stack[i, :, :] = np.array(im)
        im.close()
        return stack
#Class RegionSelector<commit_msg>Comment correction after realizing that the ImageDisplay does the transpose. Loading is essentially as expected (y in rows, x in columns)<commit_after>from PyQt5 import QtGui
from region_selector_ui import Ui_RegionSelector
import pyqtgraph as pg
import numpy as np
from PIL import Image

from utilities import RegionROI

class RegionSelector(QtGui.QMainWindow):
    """
    User interface logic to select stack regions
    """

    def __init__(self, parent=None):
        """
        Creates a new RegionSelector window
        :param parent: The window's parent
        """
        # init of ui
        QtGui.QWidget.__init__(self, parent)
        self.ui = Ui_RegionSelector()
        self.ui.setupUi(self)
        # non-ui class members
        self.__current_z = 0
        self.filename = ""
        self.currentStack = np.array([])
        self.roi_dict = {}
        self.current_z = 0
        self.current_ROI = None
        self.last_color = 0
        self.last_uid = 0
        # ui stuff settings
        self.ui.lblROIName.setText("")
        # create our view-box and image view inside stackBox
        self.stack_vbox = self.ui.stackBox.addViewBox()
        self.stack_vbox.invertY(True)  # since image display inverts the Y-axis
        self.stack_vbox.setAspectLocked(True)
        self.stack_image = pg.ImageItem()
        self.stack_vbox.addItem(self.stack_image)
        # hide menu and roi button on image views
        self.ui.ROIView.ui.roiBtn.hide()
        self.ui.ROIView.ui.menuBtn.hide()
        # connect signals
        self.ui.btnLoad.clicked.connect(self.load_click)
        self.ui.btnAddROI.clicked.connect(self.addroi_click)
        self.ui.cbRegions.currentIndexChanged.connect(self.regionNameSelChanged)
        self.ui.sldrZ.sliderMoved.connect(self.sliderZChanged)

        # for easy testing, create and display test-image stack
        test_image = np.abs(np.random.randn(10, 100, 100, 3)*50).astype(np.uint8)
        test_image[:, 50, :, :] = 255
        test_image[:, :, 50, :] = 255
        self.currentStack = test_image
        self.reset_after_load()
        self.display_slice()

    def guess_stack_type(self):
        """
        Uses simple heuristics to determine the type (2D, 3D, grayscale, color) of the current stack
        :return: Integer identifying the stack type
        """
        if self.currentStack.size == 0 or self.currentStack.ndim < 2:
            # stack invalid
            return -1
        if self.currentStack.ndim == 2:
            # grayscale 2D image
            return 0
        if self.currentStack.ndim == 3:
            if self.currentStack.shape[2] <= 4:
                # RGB(A) 2D image (or very thin stack...)
                return 1
            else:
                # greyscale 3D stack
                return 2
        if self.currentStack.ndim == 4:
            if self.currentStack.shape[3] > 4:
                # invalid color model
                return -1
            else:
                # color 3D stack
                return 3
        # no idea
        return -1

    def reset_after_load(self):
        """
        Resets user-interface elements and internal counters after a new stack has been loaded
        """
        self.roi_dict = {}
        self.current_z = 0
        self.current_ROI = None
        self.last_color = 0
        self.last_uid = 0
        self.ui.cbRegions.clear()
        self.ui.sldrZ.setMinimum(0)
        self.ui.sldrZ.setMaximum(self.NSlices - 1)
        self.ui.sldrZ.setValue(0)

    @property
    def NSlices(self):
        """
        The number of slices in the currently loaded stack
        """
        st = self.guess_stack_type()
        if st == -1:
            return 0
        if st < 2:
            return 1
        return self.currentStack.shape[0]

    @property
    def current_z(self):
        return self.__current_z

    @current_z.setter
    def current_z(self, current_z):
        if current_z >= 0:
            self.__current_z = current_z
            self.ui.lbl_z.setText(str(current_z))

    def display_slice(self):
        """
        Displays the current slice
        """
        st = self.guess_stack_type()
        if st == -1:
            print("No valid stack loaded")
            return
        if st == 0 or st == 1:
            self.stack_image.setImage(self.currentStack)
        elif st == 2 or st == 3:
            if self.NSlices > self.current_z:
                self.stack_image.setImage(self.currentStack[self.current_z])
            else:
                print("Improper slice index for stack")

    def create_default_roi(self, n_vert=6):
        """
        Create a point list for the default roi when adding a new roi
        :param n_vert: The initial number of vertices in the ROI
        :return: A list of (x,y) tuples for the vertices of the roi
        """
        st = self.guess_stack_type()
        if st == -1:
            return
        if st < 2:
            width = self.currentStack.shape[0]
            height = self.currentStack.shape[1]
        else:
            width = self.currentStack.shape[1]
            height = self.currentStack.shape[2]
        center_x = width / 2
        center_y = height / 2
        radius = min([width, height]) / 10
        point_list = []
        for i in range(n_vert):
            a = 2*np.pi / n_vert * i
            x = np.cos(a)*radius + center_x
            y = np.sin(a)*radius + center_y
            point_list.append((x, y))
        return point_list

    def next_roi_color(self):
        """
        Updates our ROI pen color cycle
        :return: The color of the next roi
        """
        self.last_color += 1
        self.last_color %= 9
        return self.last_color

    def next_roi_uid(self):
        """
        Updates the unique id of roi's
        :return: The uid of the next roi
        """
        self.last_uid += 1
        return self.last_uid

    # Signals #
    def load_click(self):
        """
        Handles event of clicking the load stack button
        """
        diag = QtGui.QFileDialog()
        fname = diag.getOpenFileName(self, "Select stack", "E:/Dropbox/2P_Data", "*.tif")[0]
        if fname is not None and fname != "":
            assert isinstance(fname, str)
            self.currentStack = self.OpenStack(fname)
            self.filename = fname
            self.reset_after_load()
            self.display_slice()
        else:
            print("No file selected")

    def addroi_click(self):
        """
        Handles event of clicking the add ROI button
        """
        name = self.ui.leNewROI.text()
        new_r = RegionROI(self.create_default_roi(), self.next_roi_uid(), name, pen=(self.next_roi_color(), 12))
        if self.current_z in self.roi_dict:
            self.roi_dict[self.current_z].append(new_r)
        else:
            self.roi_dict[self.current_z] = [new_r]
        # add the name of the created region to our list if it doesn exist yet
        found = False
        for i in range(self.ui.cbRegions.count()):
            if self.ui.cbRegions.itemText(i) == new_r.region_name:
                found = True
        if not found:
            self.ui.cbRegions.addItem(new_r.region_name)
        # add the ROI to the display
        self.stack_vbox.addItem(new_r)
        self.updateRoi(new_r)
        new_r.sigRegionChanged.connect(self.updateRoi)
        # clear the name field and combo box selection
        self.ui.leNewROI.setText("")
        self.ui.cbRegions.setCurrentIndex(-1)

    def updateRoi(self, roi: RegionROI):
        """
        Handles the event whenever an ROI is updated on-screen or programmatically
        :param roi: The roi which got updated
        """
        if roi is None:
            # clear our details image and roi indicator
            self.ui.ROIView.getImageItem().setImage(np.zeros((10,10)))
            self.ui.lblROIName.setText("")
            return
        self.current_ROI = roi
        self.ui.lblROIName.setText(roi.region_name)
        arr = roi.getArrayRegion(self.stack_image.image, img=self.stack_image)
        self.ui.ROIView.getImageItem().setImage(arr)

    def regionNameSelChanged(self, index):
        """
        Handles the event whenever a new region is selected in the combo box of previous region names
        :param index: The new index selection
        """
        if index >= 0:
            self.ui.leNewROI.setText(self.ui.cbRegions.itemText(index))

    def sliderZChanged(self, value):
        """
        Handles the event when our z-position slider was adjusted
        :param value: The new slider position
        """
        # decommission al ROI's of the current plane
        if self.current_z in self.roi_dict:
            for r in self.roi_dict[self.current_z]:
                self.stack_vbox.removeItem(r)
        # update z-plane
        self.current_z = value
        # load any existing ROI's of that plane and select the first
        if self.current_z in self.roi_dict:
            for i, r in enumerate(self.roi_dict[self.current_z]):
                self.stack_vbox.addItem(r)
                if i == 0:
                    self.updateRoi(r)
        else:
            self.updateRoi(None)
        # display the new slice
        self.display_slice()

    @staticmethod
    def OpenStack(filename):
        """
        Load image stack from tiff file
        :param filename: The name of the tiff file
        :return: A (z, x, y) numpy array
        """
        im = Image.open(filename)
        # NOTE: In the follwoing im.size[1] is the height of the image (0-ymax) and im.size[0]
        # is the width (0-xmax). Increasing y are along rows, increasing x along columns
        if im.mode == 'RGB':
            stack = np.empty((im.n_frames, im.size[1], im.size[0], 3), dtype=np.uint8)
        else:
            stack = np.empty((im.n_frames, im.size[1], im.size[0]), dtype=np.uint8)
        # loop over frames and assign
        for i in range(im.n_frames):
            im.seek(i)
            stack[i, :, :] = np.array(im)
        im.close()
        return stack
#Class RegionSelector<|endoftext|>


<commit_before>from obspy.segy.core import readSEGY
from las import LASReader

import fnmatch

import os
from itertools import product

from matplotlib import pyplot as plt, gridspec
import numpy as np
from scipy.interpolate import griddata

import rasterio
from fiona import collection
from shapely.geometry import shape, mapping, LineString, Point
from shapely.ops import unary_union
from shapely.prepared import prep

from plot_lib import uberPlot
from lithology.lithology import intervals_from_las3_string

import pyproj as pp

class transectContainer():
    """
    Main driver class for generating transect plots. Builds and
    accesses all other plot containers

    usage:
    tc = transectContainer(transect_dir, seismic_dir,
                           elevation_raster, las_dir,
                           extents)

    @param transect_dir: Directory containing shape files of the
                         transects.
    @param seismic_dir: Directory containing the shape files of
                        the SEGY trace headers.
    @param elevation_raster: Raster file of the entire elevation
                             profile.
    @param las_dir: Directory containing shape files for well log
                    headers.

    @param extents: X and depth limits of the plot (X0,X1, Z0, Z1)

    @returns a transectContainer object.

    tc.plot() to generate transect plots
    """

    def __init__(self, transect_dir, seismic_dir,
                 elevation_raster, las_dir,
                 bedrock_dir, striplog_dir,
                 extents):

        self.extents = extents

        # make the data objects
        self.seismic = seismicContainer(seismic_dir)
        self.elevation = elevationContainer(elevation_raster)
        self.log = lasContainer(las_dir)
        self.bedrock = bedrockContainer(bedrock_dir)
        self.striplog = striplogContainer(striplog_dir)

        # Place holder for em/gravity etc...
        self.dummy = dummyContainer()

        # initialize transect list
        self.data = []

        # Read in all shape files
        for f in os.listdir(transect_dir):

            if not f.endswith(".shp"):
                continue

            for line in collection(os.path.join(transect_dir,f)):

                self.data.append(shape(line['geometry']))

    def plot(self):
        """
        Generates transect plots for each transect line
        """

        for transect in self.data:

            # Set the extent to the length? Or keep them all the same?
            self.extents[1] = transect.length

            # update the containers
            self.seismic.update(transect)
            self.log.update(transect)
            self.elevation.update(transect)
            self.bedrock.update(transect)
            self.striplog.update(transect)
            self.dummy.update(transect)

            uberPlot(transect, self.seismic, self.elevation,
                     self.log, self.bedrock,
                     self.striplog, self.dummy,
                     self.extents)
            plt.show()

class seismicContainer():
    """
    Class for reading and plotting seismic data.

    USAGE:

    seis = seismicContainer(seis_dir)

    @param seis_dir: Input directory containing seismic shapefiles.
                     The shapefiles contain points corresponding to
                     UTM trace locations with fields segyfile,
                     trace.

    @returns a seismicContainer object
    """
    def __init__(self, seis_dir):

        self.lookup = {}    # Look up table for points: segy/trace
        self.data = []      # plotting data
        self.coords = []    # transect coords of plot data
        self.buffer = 300.0 # [m] of buffer for transect association

        for f in fnmatch.filter(os.listdir(seis_dir), '*.shp'):

            shapefile = os.path.join(seis_dir, f)

            with collection(shapefile, "r") as traces:

                for trace in traces:

                    # Map a point object to trace properties
                    # (segyfile, trace)
                    self.lookup[shape(trace["geometry"])] =\
                                trace["properties"]

    def update(self, transect):
        """
        Updates the container data to traces that intersect
        the transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # Preprocessing
        prepared = prep(transect.buffer(self.buffer))

        # Get intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # Reset data
        self.data = []
        self.coords = []

        # Lookup for grouping traces into segy files
        file_lookup = {}
        for point in points:

            meta = self.lookup[point]

            if(meta["segyfile"] in file_lookup):

                # project onto the transect line
                proj_d = transect.project(point)

                if proj_d:
                    file_lookup[meta["segyfile"]]["pos"].append(proj_d)
                    file_lookup[meta["segyfile"]]["trace"].append(meta["trace"])

            else: ## Make a new dict entry

                file_lookup[meta["segyfile"]] = {}
                file_lookup[meta["segyfile"]]["trace"] = \
                  [meta["trace"]]

                proj_d = transect.project(point)
                file_lookup[meta["segyfile"]]["pos"] = [proj_d]

        # Read in the chunks from the segy file
        for segyfile in file_lookup.keys():

            segy = readSEGY(segyfile)
            traces = file_lookup[segyfile]["trace"]
            coords = file_lookup[segyfile]["pos"]

            # sort the traces to be in order
            idx = sorted(range(len(traces)), key=lambda k: traces[k])

            self.data.append(np.transpose(np.array(
                [segy.traces[traces[i]].data for i in idx])))

            self.coords.append(np.array([coords[i] for i in idx]))

class lasContainer():
    """
    Container for managing and plotting LAS data.

    usage:

    lc = lasContainer(las_dir)

    @param las_dir: Directory shape files of LAS headers.

    @returns an lasContainer object/
    """
    def __init__(self, las_dir):

        # Read in the shape file
        self.lookup = {}   # maps points to LAS filenames
        self.data = []     # plot data
        self.coords = []   # transect coords of plot data
        self.buffer = 300  # [m] buffer for transect association
        self.log_lookup  ={} # look up for log ids

        for root, dirs, files in os.walk(las_dir):

            try:
                shapefile = \
                  os.path.join(root,fnmatch.filter(files,'*.shp')[0])
            except:
                continue

            with collection(shapefile, "r") as logs:

                for log in logs:

                    # add to the lookup table
                    self.lookup[shape(log["geometry"])] =\
                        log["properties"]

    def update(self, transect):
        """
        Updates the container data to wells that intersect the
        transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # Preprocess
        prepared = prep(transect.buffer(self.buffer))

        # Get the intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # reset the data
        self.data = []
        self.coords = []
        self.log_lookup = {}

        for point in points:

            meta = self.lookup[point]

            name = meta["name"]
            filename = os.path.join('data', 'wells', name,
                                    'wireline_log', name +
                                    '_out.LAS')

            # Write out an error log?
            if not os.path.exists(filename): continue

            # store the transecting data
            self.data.append(LASReader(filename, null_subs=np.nan))
            self.coords.append(transect.project(point))
            self.log_lookup[name] = self.data[-1]

    def get(self, log_id):
        """
        Returns data corresponding to log_id
        """

        return self.log_lookup.get(log_id)

class elevationContainer():
    """
    Container for managing and plotting elevation data.

    usage:
    ec = elevationContainer(elevation_raster)

    @param elevation_raster: Raster file containing the elevation
                             profile of the area of interest.

    @returns an elevationContainer object.
    """

    def __init__(self, elevation_file):

        # entire data set and grid
        self.elevation_profile = []
        self.elevation_grid = []

        # plotting data and transect coordinates
        self.data = []
        self.coords = []

        # decimation factor
        decimate = 1

        with rasterio.drivers(CPL_DEBUG=True):
            with rasterio.open(elevation_file) as src:
                self.elevation_profile = src.read()[0,0:-1:decimate,
                                                    0:-1:decimate]

                # Get as lat/lon using the affine coordinate
                # transform
                lat = np.arange(self.elevation_profile.shape[1]) \
                  *src.affine[0]*decimate + src.affine[2]

                lon = np.arange(self.elevation_profile.shape[0])\
                   * src.affine[4]*decimate + src.affine[5]

                wgs_grid = np.meshgrid(lat,lon)

                # Maybe these should be config params and not hard
                # coded?
                ll_wgs84 = pp.Proj("+init=EPSG:4269")
                utm_nad83 = pp.Proj("+init=EPSG:26920")

                self.elevation_grid = pp.transform(ll_wgs84,
                                                   utm_nad83,
                                                   wgs_grid[0],
                                                   wgs_grid[1])

    def update(self, transect):
        """
        Updates the container data to a profile that intersect the
        transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # transect coords need to be upsampled
        nsamples = 100

        self.coords = np.zeros(nsamples)
        self.data = np.zeros(nsamples)

        for i,n in enumerate(np.linspace(0,transect.length,
                                         nsamples)):

            # interpolate along the transect
            x,y = transect.interpolate(n).xy

            # Get the closest elevation points
            xi = np.abs(self.elevation_grid[0][0,:] - x).argmin()
            yi = np.abs(self.elevation_grid[1][:,0] - y).argmin()

            self.data[i] = self.elevation_profile[yi,xi]

            # add the distance to the coordinates
            self.coords[i] = n

    def plot(self, extents, bedrock):
        """
        Plots the elevation profile.
        @uses elevation_plot

        @param bedrock: bedrockContainer object
        @param extents: Plot extents (x0,x1,z0,z1). Only x0 and x1
                        are used.
        """

        elevation_plot(self, bedrock,
                       [extents[0], extents[1]],
                       np.amax(self.elevation_profile))

class dummyContainer():
    # Random number generator for testing

    def __init__(self):

        self.data = []
        self.coords = []

    def update(self, transect):

        self.coords = np.linspace(0, transect.length, 1000)
        self.data = np.random.randn(self.coords.size)

    def plot(self, extents, xticks):

        plt.plot(self.coords, self.data)
        plt.yticks([-4,0,4])

        plt.xticks(xticks, [])

        plt.ylabel("Anomaly [mGal]", fontsize=8)
        plt.tick_params(axis='y', which="major", labelsize=8)

        plt.grid(True)
        plt.xlim((extents[0], extents[1]))
        plt.ylim((-4,4))

class bedrockContainer():

    def __init__(self, bedrock_dir):

        self.lookup = {}

        self.data = []
        self.coords = []

        self.buffer = 1000 #[m]

        # Read in all shape files
        for f in os.listdir(bedrock_dir):

            if not f.endswith(".shp"):
                continue

            for line in collection(os.path.join(bedrock_dir,f)):

                self.lookup[(shape(line['geometry']))] = \
                            line['properties']

    def update(self, transect):

        points = [Point(xy) for xy in transect.coords]

        self.data = []
        self.coords = []

        for polygon, properties in self.lookup.items():

            for p in filter(polygon.contains, points):

                self.data.append(properties)
                self.coords.append(transect.project(p))

        # sort the data to be in order
        idx = sorted(range(len(self.coords)),
                     key=lambda k: self.coords[k])

        self.data = [self.data[i] for i in idx]
        self.coords = np.array([self.coords[i] for i in idx])

class striplogContainer(lasContainer):

    def update(self, transect):

        # Preprocess
        prepared = prep(transect.buffer(self.buffer))

        # Get the intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # reset the data
        self.data = []
        self.coords = []
        self.log_lookup = {}

        for point in points:

            meta = self.lookup[point]

            name = meta["name"]
            filename = os.path.join('data', 'wells', name,
                                    'lithology_log', name +
                                    '_striplog.las')

            # Write out an error log?
            if not os.path.exists(filename): continue

            # store the transecting data
            data = LASReader(filename, null_subs=np.nan,
                             unknown_as_other=True)

            self.data.append(intervals_from_las3_string(data.other))
            self.coords.append(transect.project(point))
            self.log_lookup[name] = self.data[-1]

    def plot(self):

        plot_striplog(self.data)

<commit_msg>prints for ev<commit_after>from obspy.segy.core import readSEGY
from las import LASReader

import fnmatch

import os
from itertools import product

from matplotlib import pyplot as plt, gridspec
import numpy as np
from scipy.interpolate import griddata

import rasterio
from fiona import collection
from shapely.geometry import shape, mapping, LineString, Point
from shapely.ops import unary_union
from shapely.prepared import prep

from plot_lib import uberPlot
from lithology.lithology import intervals_from_las3_string

import pyproj as pp

class transectContainer():
    """
    Main driver class for generating transect plots. Builds and
    accesses all other plot containers

    usage:
    tc = transectContainer(transect_dir, seismic_dir,
                           elevation_raster, las_dir,
                           extents)

    @param transect_dir: Directory containing shape files of the
                         transects.
    @param seismic_dir: Directory containing the shape files of
                        the SEGY trace headers.
    @param elevation_raster: Raster file of the entire elevation
                             profile.
    @param las_dir: Directory containing shape files for well log
                    headers.

    @param extents: X and depth limits of the plot (X0,X1, Z0, Z1)

    @returns a transectContainer object.

    tc.plot() to generate transect plots
    """

    def __init__(self, transect_dir, seismic_dir,
                 elevation_raster, las_dir,
                 bedrock_dir, striplog_dir,
                 extents):

        self.extents = extents

        # make the data objects
        self.seismic = seismicContainer(seismic_dir)
        self.elevation = elevationContainer(elevation_raster)
        self.log = lasContainer(las_dir)
        self.bedrock = bedrockContainer(bedrock_dir)
        self.striplog = striplogContainer(striplog_dir)

        # Place holder for em/gravity etc...
        self.dummy = dummyContainer()

        # initialize transect list
        self.data = []

        # Read in all shape files
        for f in os.listdir(transect_dir):

            if not f.endswith(".shp"):
                continue

            for line in collection(os.path.join(transect_dir,f)):

                self.data.append(shape(line['geometry']))

    def plot(self):
        """
        Generates transect plots for each transect line
        """

        for transect in self.data:

            # Set the extent to the length? Or keep them all the same?
            self.extents[1] = transect.length

            # update the containers
            self.seismic.update(transect)
            self.log.update(transect)
            self.elevation.update(transect)
            self.bedrock.update(transect)
            self.striplog.update(transect)
            self.dummy.update(transect)

            uberPlot(transect, self.seismic, self.elevation,
                     self.log, self.bedrock,
                     self.striplog, self.dummy,
                     self.extents)
            plt.show()

class seismicContainer():
    """
    Class for reading and plotting seismic data.

    USAGE:

    seis = seismicContainer(seis_dir)

    @param seis_dir: Input directory containing seismic shapefiles.
                     The shapefiles contain points corresponding to
                     UTM trace locations with fields segyfile,
                     trace.

    @returns a seismicContainer object
    """
    def __init__(self, seis_dir):

        self.lookup = {}    # Look up table for points: segy/trace
        self.data = []      # plotting data
        self.coords = []    # transect coords of plot data
        self.buffer = 300.0 # [m] of buffer for transect association

        for f in fnmatch.filter(os.listdir(seis_dir), '*.shp'):

            print("Found shape files for seismic")
            shapefile = os.path.join(seis_dir, f)

            with collection(shapefile, "r") as traces:

                for trace in traces:

                    print("Found seismic trace locations")
                    # Map a point object to trace properties
                    # (segyfile, trace)
                    self.lookup[shape(trace["geometry"])] =\
                                trace["properties"]

    def update(self, transect):
        """
        Updates the container data to traces that intersect
        the transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # Preprocessing
        prepared = prep(transect.buffer(self.buffer))

        # Get intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # Reset data
        self.data = []
        self.coords = []

        # Lookup for grouping traces into segy files
        file_lookup = {}
        for point in points:

            print("updating seismic transect")
            meta = self.lookup[point]

            if(meta["segyfile"] in file_lookup):

                # project onto the transect line
                proj_d = transect.project(point)

                if proj_d:
                    file_lookup[meta["segyfile"]]["pos"].append(proj_d)
                    file_lookup[meta["segyfile"]]["trace"].append(meta["trace"])

            else: ## Make a new dict entry

                file_lookup[meta["segyfile"]] = {}
                file_lookup[meta["segyfile"]]["trace"] = \
                  [meta["trace"]]

                proj_d = transect.project(point)
                file_lookup[meta["segyfile"]]["pos"] = [proj_d]

        # Read in the chunks from the segy file
        for segyfile in file_lookup.keys():

            segy = readSEGY(segyfile)
            traces = file_lookup[segyfile]["trace"]
            coords = file_lookup[segyfile]["pos"]

            print("Read in seismic data")
            # sort the traces to be in order
            idx = sorted(range(len(traces)), key=lambda k: traces[k])

            self.data.append(np.transpose(np.array(
                [segy.traces[traces[i]].data for i in idx])))

            self.coords.append(np.array([coords[i] for i in idx]))

class lasContainer():
    """
    Container for managing and plotting LAS data.

    usage:

    lc = lasContainer(las_dir)

    @param las_dir: Directory shape files of LAS headers.

    @returns an lasContainer object/
    """
    def __init__(self, las_dir):

        # Read in the shape file
        self.lookup = {}   # maps points to LAS filenames
        self.data = []     # plot data
        self.coords = []   # transect coords of plot data
        self.buffer = 300  # [m] buffer for transect association
        self.log_lookup  ={} # look up for log ids

        for root, dirs, files in os.walk(las_dir):

            try:
                shapefile = \
                  os.path.join(root,fnmatch.filter(files,'*.shp')[0])
            except:
                continue

            with collection(shapefile, "r") as logs:

                for log in logs:

                    # add to the lookup table
                    self.lookup[shape(log["geometry"])] =\
                        log["properties"]

    def update(self, transect):
        """
        Updates the container data to wells that intersect the
        transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # Preprocess
        prepared = prep(transect.buffer(self.buffer))

        # Get the intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # reset the data
        self.data = []
        self.coords = []
        self.log_lookup = {}

        for point in points:

            meta = self.lookup[point]

            name = meta["name"]
            filename = os.path.join('data', 'wells', name,
                                    'wireline_log', name +
                                    '_out.LAS')

            # Write out an error log?
            if not os.path.exists(filename): continue

            # store the transecting data
            self.data.append(LASReader(filename, null_subs=np.nan))
            self.coords.append(transect.project(point))
            self.log_lookup[name] = self.data[-1]

    def get(self, log_id):
        """
        Returns data corresponding to log_id
        """

        return self.log_lookup.get(log_id)

class elevationContainer():
    """
    Container for managing and plotting elevation data.

    usage:
    ec = elevationContainer(elevation_raster)

    @param elevation_raster: Raster file containing the elevation
                             profile of the area of interest.

    @returns an elevationContainer object.
    """

    def __init__(self, elevation_file):

        # entire data set and grid
        self.elevation_profile = []
        self.elevation_grid = []

        # plotting data and transect coordinates
        self.data = []
        self.coords = []

        # decimation factor
        decimate = 1

        with rasterio.drivers(CPL_DEBUG=True):
            with rasterio.open(elevation_file) as src:
                self.elevation_profile = src.read()[0,0:-1:decimate,
                                                    0:-1:decimate]

                # Get as lat/lon using the affine coordinate
                # transform
                lat = np.arange(self.elevation_profile.shape[1]) \
                  *src.affine[0]*decimate + src.affine[2]

                lon = np.arange(self.elevation_profile.shape[0])\
                   * src.affine[4]*decimate + src.affine[5]

                wgs_grid = np.meshgrid(lat,lon)

                # Maybe these should be config params and not hard
                # coded?
                ll_wgs84 = pp.Proj("+init=EPSG:4269")
                utm_nad83 = pp.Proj("+init=EPSG:26920")

                self.elevation_grid = pp.transform(ll_wgs84,
                                                   utm_nad83,
                                                   wgs_grid[0],
                                                   wgs_grid[1])

    def update(self, transect):
        """
        Updates the container data to a profile that intersect the
        transect line.

        @param transect: A transect line as a shapely LineString
                         object.
        """

        # transect coords need to be upsampled
        nsamples = 100

        self.coords = np.zeros(nsamples)
        self.data = np.zeros(nsamples)

        for i,n in enumerate(np.linspace(0,transect.length,
                                         nsamples)):

            # interpolate along the transect
            x,y = transect.interpolate(n).xy

            # Get the closest elevation points
            xi = np.abs(self.elevation_grid[0][0,:] - x).argmin()
            yi = np.abs(self.elevation_grid[1][:,0] - y).argmin()

            self.data[i] = self.elevation_profile[yi,xi]

            # add the distance to the coordinates
            self.coords[i] = n

    def plot(self, extents, bedrock):
        """
        Plots the elevation profile.
        @uses elevation_plot

        @param bedrock: bedrockContainer object
        @param extents: Plot extents (x0,x1,z0,z1). Only x0 and x1
                        are used.
        """

        elevation_plot(self, bedrock,
                       [extents[0], extents[1]],
                       np.amax(self.elevation_profile))

class dummyContainer():
    # Random number generator for testing

    def __init__(self):

        self.data = []
        self.coords = []

    def update(self, transect):

        self.coords = np.linspace(0, transect.length, 1000)
        self.data = np.random.randn(self.coords.size)

    def plot(self, extents, xticks):

        plt.plot(self.coords, self.data)
        plt.yticks([-4,0,4])

        plt.xticks(xticks, [])

        plt.ylabel("Anomaly [mGal]", fontsize=8)
        plt.tick_params(axis='y', which="major", labelsize=8)

        plt.grid(True)
        plt.xlim((extents[0], extents[1]))
        plt.ylim((-4,4))

class bedrockContainer():

    def __init__(self, bedrock_dir):

        self.lookup = {}

        self.data = []
        self.coords = []

        self.buffer = 1000 #[m]

        # Read in all shape files
        for f in os.listdir(bedrock_dir):

            if not f.endswith(".shp"):
                continue

            for line in collection(os.path.join(bedrock_dir,f)):

                self.lookup[(shape(line['geometry']))] = \
                            line['properties']

    def update(self, transect):

        points = [Point(xy) for xy in transect.coords]

        self.data = []
        self.coords = []

        for polygon, properties in self.lookup.items():

            for p in filter(polygon.contains, points):

                self.data.append(properties)
                self.coords.append(transect.project(p))

        # sort the data to be in order
        idx = sorted(range(len(self.coords)),
                     key=lambda k: self.coords[k])

        self.data = [self.data[i] for i in idx]
        self.coords = np.array([self.coords[i] for i in idx])

class striplogContainer(lasContainer):

    def update(self, transect):

        # Preprocess
        prepared = prep(transect.buffer(self.buffer))

        # Get the intersecting points
        points = filter(prepared.contains, self.lookup.keys())

        # reset the data
        self.data = []
        self.coords = []
        self.log_lookup = {}

        for point in points:

            meta = self.lookup[point]

            name = meta["name"]
            filename = os.path.join('data', 'wells', name,
                                    'lithology_log', name +
                                    '_striplog.las')

            # Write out an error log?
            if not os.path.exists(filename): continue

            # store the transecting data
            data = LASReader(filename, null_subs=np.nan,
                             unknown_as_other=True)

            self.data.append(intervals_from_las3_string(data.other))
            self.coords.append(transect.project(point))
            self.log_lookup[name] = self.data[-1]

    def plot(self):

        plot_striplog(self.data)

<|endoftext|>


<commit_before>from CouncilTag import settings
from datetime import datetime, timedelta
import logging
import requests
import os
import bcrypt
import base64
import pytz
import googlemaps
import boto3
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.application import MIMEApplication
log = logging.Logger(__name__)
ses_client = boto3.client('ses', aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
                          aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
                          region_name=os.environ["AWS_REGION"])

def verify_recaptcha(token):
    r = requests.post('https://www.google.com/recaptcha/api/siteverify', data={
                      'secret': os.environ["RECAPTCHAKEY"], 'response': token})
    response = r.json()
    return response['success']

def array_of_ordereddict_to_list_of_names(tags_ordereddict_array):
    """
    Serializers have a funny organization that isn't helpful in making further queries
    Here we take the list of ordered dictionaries (id: x, name: y) and pull out the name only
    and put that in a names list to return
    """
    names = []
    length = len(list(tags_ordereddict_array))
    for i in range(length):
        names.append(tags_ordereddict_array[i]["name"])
    return names

def check_auth_code(plain_code, hashed):
    dec = bcrypt.hashpw(plain_code.encode('utf-8'),
                        hashed.encode('utf-8')).decode('utf-8')
    if dec == hashed:
        return True
    return False

def calculateTallies(messages_qs):
    pro = 0
    con = 0
    more_info = 0
    home_owner = 0
    business_owner = 0
    resident = 0
    works = 0
    school = 0
    child_school = 0
    total = 0
    for message in messages_qs:
        if message.authcode != None:
            continue
        if message.pro == 0:
            con += 1
        elif message.pro == 1:
            pro += 1
        else:
            more_info += 1
        if message.home_owner:
            home_owner += 1
        if message.business_owner:
            business_owner += 1
        if message.resident:
            resident += 1
        if message.works:
            works += 1
        if message.school:
            school += 1
        if message.child_school:
            child_school += 1
        total += 1
    return {"home_owner": home_owner, "business_owner": business_owner,
            "resident": resident, "works": works, "school": school,
            "child_school": child_school, "pro": pro, "con": con, "more_info": more_info, "total": total}

def getLocationBasedDate(timestamp, cutoff_days_offset, cutoff_hour, cutoff_minute, location_tz):
    """
    @timestamp a UTC timestamp
    @cutoff_dats_offset +/- integer days from now that should be checking date
    @cutoff_hours integer hours for location that should be set
    @cutoff_minutes integer minutes for locaiton that should be set
    @location_tz string timezone for location where meeting takes place
    """
    tz = pytz.timezone(location_tz)
    dt = datetime.utcfromtimestamp(timestamp)
    dt = dt.astimezone(tz)
    if cutoff_days_offset is not None:
        dt = dt + timedelta(days=cutoff_days_offset)
    if cutoff_hour is not None:
        dt = dt.replace(hour=cutoff_hour, minute=cutoff_minute)
    return dt

def isCommentAllowed(timestamp, cutoff_days_offset, cutoff_hours, cutoff_minutes, location_tz):
    dt = getLocationBasedDate(
        timestamp, cutoff_days_offset, cutoff_hours, cutoff_minutes, location_tz)
    now = datetime.now().astimezone(tz=pytz.timezone(location_tz))
    if (now > dt):
        return False
    return True

def send_mail(mail_message):
    if type(mail_message["user"]) is dict:
        to_email = mail_message["user"]["email"]
    else:
        to_email = mail_message["user"].email
    multipart_content_subtype = 'mixed'
    msg = MIMEMultipart(multipart_content_subtype)
    msg['Subject'] = mail_message["subject"]
    msg['To'] = to_email
    msg['From'] = 'do-not-reply@engage.town'
    part = MIMEText(mail_message['content'], 'html')
    msg.attach(part)
    if "attachment_file_path" in mail_message:
        with open(mail_message["attachment_file_path"], 'rb') as f:
            part = MIMEApplication(f.read(), _subtype='pdf')
            part.add_header('Content-Disposition', 'attachment',
                            filename=mail_message['attachment_file_name'])
            msg.attach(part)
    response = ses_client.send_raw_email(
        Source="engage team <do-not-reply@engage.town>",
        Destinations=[to_email],
        RawMessage={'Data': msg.as_string()})
    if response['MessageId'] is not None:
        return True
    else:
        log.error("Could not send an email from {} to {} about {}".format(from_email,
                                                                          to_email, subject))
        log.error(response.body)
        log.error(response.status_code)
        return False
<commit_msg>slight fix<commit_after>from CouncilTag import settings
from datetime import datetime, timedelta
import logging
import requests
import os
import bcrypt
import base64
import pytz
import googlemaps
import boto3
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.application import MIMEApplication
log = logging.Logger(__name__)
ses_client = boto3.client('ses', aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
                          aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
                          region_name=os.environ["AWS_REGION"])

def verify_recaptcha(token):
    r = requests.post('https://www.google.com/recaptcha/api/siteverify', data={
                      'secret': os.environ["RECAPTCHAKEY"], 'response': token})
    response = r.json()
    return response['success']

def array_of_ordereddict_to_list_of_names(tags_ordereddict_array):
    """
    Serializers have a funny organization that isn't helpful in making further queries
    Here we take the list of ordered dictionaries (id: x, name: y) and pull out the name only
    and put that in a names list to return
    """
    names = []
    length = len(list(tags_ordereddict_array))
    for i in range(length):
        names.append(tags_ordereddict_array[i]["name"])
    return names

def check_auth_code(plain_code, hashed):
    dec = bcrypt.hashpw(plain_code.encode('utf-8'),
                        hashed.encode('utf-8')).decode('utf-8')
    if dec == hashed:
        return True
    return False

def calculateTallies(messages_qs):
    pro = 0
    con = 0
    more_info = 0
    home_owner = 0
    business_owner = 0
    resident = 0
    works = 0
    school = 0
    child_school = 0
    total = 0
    for message in messages_qs:
        if message.authcode != None:
            continue
        if message.pro == 0:
            con += 1
        elif message.pro == 1:
            pro += 1
        else:
            more_info += 1
        if message.home_owner:
            home_owner += 1
        if message.business_owner:
            business_owner += 1
        if message.resident:
            resident += 1
        if message.works:
            works += 1
        if message.school:
            school += 1
        if message.child_school:
            child_school += 1
        total += 1
    return {"home_owner": home_owner, "business_owner": business_owner,
            "resident": resident, "works": works, "school": school,
            "child_school": child_school, "pro": pro, "con": con, "more_info": more_info, "total": total}

def getLocationBasedDate(timestamp, cutoff_days_offset, cutoff_hour, cutoff_minute, location_tz):
    """
    @timestamp a UTC timestamp
    @cutoff_dats_offset +/- integer days from now that should be checking date
    @cutoff_hours integer hours for location that should be set
    @cutoff_minutes integer minutes for locaiton that should be set
    @location_tz string timezone for location where meeting takes place
    """
    tz = pytz.timezone(location_tz)
    dt = datetime.utcfromtimestamp(timestamp)
    dt = dt.astimezone(tz)
    if cutoff_days_offset is not None:
        dt = dt + timedelta(days=cutoff_days_offset)
    if cutoff_hour is not None:
        dt = dt.replace(hour=cutoff_hour, minute=cutoff_minute)
    return dt

def isCommentAllowed(timestamp, cutoff_days_offset, cutoff_hours, cutoff_minutes, location_tz):
    dt = getLocationBasedDate(
        timestamp, cutoff_days_offset, cutoff_hours, cutoff_minutes, location_tz)
    now = datetime.now().astimezone(tz=pytz.timezone(location_tz))
    if (now > dt):
        return False
    return True

def send_mail(mail_message):
    if type(mail_message["user"]) is dict:
        to_email = mail_message["user"]["email"]
    else:
        to_email = mail_message["user"].email
    multipart_content_subtype = 'mixed'
    msg = MIMEMultipart(multipart_content_subtype)
    msg['Subject'] = mail_message["subject"]
    msg['To'] = to_email
    msg['From'] = 'do-not-reply@engage.town'
    part = MIMEText(mail_message['content'], 'html')
    msg.attach(part)
    if "attachment_file_path" in mail_message:
        with open(mail_message["attachment_file_path"], 'rb') as f:
            part = MIMEApplication(f.read(), _subtype='pdf')
            part.add_header('Content-Disposition', 'attachment',
                            filename=mail_message['attachment_file_name'])
            msg.attach(part)
    response = ses_client.send_raw_email(
        Source="engage team <do-not-reply@engage.town>",
        Destinations=[to_email],
        RawMessage={'Data': msg.as_string()})
    if response['MessageId'] is not None:
        return True
    else:
        log.error("Could not send an email from {} to {} about {}".format("do-not-reply@engage.town",
                                                                          to_email, mail_message['Subject']))
        log.error(response.body)
        log.error(response.status_code)
        return False
<|endoftext|>


<commit_before>import time
import twitter
import json

def oauth_login():
    # credentials for OAuth
    CONSUMER_KEY = c_k
    CONSUMER_SECRET = c_s
    OAUTH_TOKEN = a_t
    OAUTH_TOKEN_SECRET = a_t_s
    # Creating the authentification
    auth = twitter.oauth.OAuth(OAUTH_TOKEN,
                               OAUTH_TOKEN_SECRET,
                               CONSUMER_KEY,
                               CONSUMER_SECRET)
    # Twitter instance
    twitter_api = twitter.Twitter(auth=auth)
    return twitter_api

dataFile = open('data', 'w')
#seedWords = ['eating now', 'watching now', 'making now ', 'heading now', 'having now', 'doing now', 'working now']

c_k = '1YjJerf1D5rRvriEzRstKw'
c_s = 'gupHGRRSvQoccMDqbw3slw5KgfV6hYULvjWbr9i2nR4'
a_t = '28816243-o41jYj1BmM40ARLaJzdpJUsosBQhfgah3UJBlrgEU'
a_t_s = 'ZlHbMhluZv0yBlG8cIkMAoFX8XK0pxpFn39OKAqTaQ'

#size = len(seedWords)
index = 0
requestLimit = 180
requestNum = 0
idSet = set()
while True:
    twitter_api = oauth_login()
    requestNum += 1
    if requestNum > requestLimit:
        print 'Wait for 15 mins...'
        time.sleep(900)
        requestNum = 1
    try:
        #print "Collecting Tweets: " + str(seedWords[index])
        print "Collecting Tweets..."
        tweets = twitter_api.search.tweets(language='en', count=100, include_entities=True,
                                           q='"now"')
        #index += 1
        #if index >= size:
        #    index = 0
    except Exception as e:  # take care of errors
        print 'API ERROR: ' + str(e)
        #if index >= size:
        #    index = 0
        continue
    for tweet in tweets['statuses']:
        if tweet['id'] not in idSet:
            idSet.add(tweet['id'])
            temp = {}
            temp['text'] = tweet['text']
            temp['id'] = tweet['id']
            temp['entities'] = tweet['entities']
            temp['geo'] = tweet['geo']
            temp['place'] = tweet['place']
            temp['source'] = tweet['source']
            temp['time'] = tweet['created_at']
            temp['user_lang'] = tweet['user']['lang']
            temp['user_location'] = tweet['user']['location']
            temp['user_timezone'] = tweet['user']['time_zone']
            dataFile.write(json.dumps(temp))
            dataFile.write('\n')

dataFile.close()<commit_msg>update collector<commit_after>import time
import twitter
import json

def oauth_login():
    # credentials for OAuth
    CONSUMER_KEY = c_k
    CONSUMER_SECRET = c_s
    OAUTH_TOKEN = a_t
    OAUTH_TOKEN_SECRET = a_t_s
    # Creating the authentification
    auth = twitter.oauth.OAuth(OAUTH_TOKEN,
                               OAUTH_TOKEN_SECRET,
                               CONSUMER_KEY,
                               CONSUMER_SECRET)
    # Twitter instance
    twitter_api = twitter.Twitter(auth=auth)
    return twitter_api

dataFile = open('data', 'a')
userFile = open('user', 'a')
#seedWords = ['eating now', 'watching now', 'making now ', 'heading now', 'having now', 'doing now', 'working now']

c_k = '1YjJerf1D5rRvriEzRstKw'
c_s = 'gupHGRRSvQoccMDqbw3slw5KgfV6hYULvjWbr9i2nR4'
a_t = '28816243-o41jYj1BmM40ARLaJzdpJUsosBQhfgah3UJBlrgEU'
a_t_s = 'ZlHbMhluZv0yBlG8cIkMAoFX8XK0pxpFn39OKAqTaQ'

#size = len(seedWords)
index = 0
requestLimit = 180
requestNum = 0
idSet = set()
userIDSet = set()
while True:
    twitter_api = oauth_login()
    requestNum += 1
    if requestNum > requestLimit:
        print 'Wait for 15 mins...'
        time.sleep(900)
        requestNum = 1
    try:
        #print "Collecting Tweets: " + str(seedWords[index])
        print "Collecting Tweets, request num: "+str(requestNum)
        tweets = twitter_api.search.tweets(language='en', count=100, include_entities=True, q='e')
        #index += 1
        #if index >= size:
        #    index = 0
    except Exception as e:  # take care of errors
        print 'API ERROR: ' + str(e)
        #if index >= size:
        #    index = 0
        continue
    for tweet in tweets['statuses']:
        if tweet['id'] not in idSet:
            idSet.add(tweet['id'])
            temp = {}
            temp['text'] = tweet['text']
            temp['id'] = tweet['id']
            temp['entities'] = tweet['entities']
            temp['geo'] = tweet['geo']
            temp['place'] = tweet['place']
            temp['source'] = tweet['source']
            temp['time'] = tweet['created_at']
            temp['user_lang'] = tweet['user']['lang']
            temp['user_location'] = tweet['user']['location']
            temp['user_timezone'] = tweet['user']['time_zone']
            temp['user_name'] = tweet['user']['screen_name']
            temp['user_id'] = tweet['user']['id']
            if temp['user_id'] not in userIDSet:
                userIDSet.add(temp['user_id'])
                userFile.write(str(temp['user_id'])+'\t'+temp['user_name']+'\n')
            dataFile.write(json.dumps(temp))
            dataFile.write('\n')

dataFile.close()
userFile.close()<|endoftext|>


<commit_before>from __future__ import print_function

from random import choice
from itertools import chain
from accumulation_tree import AccumulationTree
import pyudorandom

def _centroid_count(centroid):
    # this is defined at the top level because
    # pyspark needs to pickle these functions properly
    # See https://github.com/CamDavidsonPilon/tdigest/issues/42
    return centroid.count

class Centroid(object):

    def __init__(self, mean, count):
        self.mean = float(mean)
        self.count = float(count)

    def __repr__(self):
        return """<Centroid: mean=%.8f, count=%d>""" % (self.mean, self.count)

    def __eq__(self, other):
        return self.mean == other.mean and self.count == other.count

    def update(self, x, weight):
        self.count += weight
        self.mean += weight * (x - self.mean) / self.count
        return

class TDigest(object):

    def __init__(self, delta=0.01, K=25):

        self.C = AccumulationTree(_centroid_count)
        self.n = 0
        self.delta = delta
        self.K = K

    def __add__(self, other_digest):
        data = list(chain(self.C.values(), other_digest.C.values()))
        new_digest = TDigest(self.delta, self.K)

        if len(data) > 0:
            for c in pyudorandom.items(data):
                new_digest.update(c.mean, c.count)

        return new_digest

    def __len__(self):
        return len(self.C)

    def __repr__(self):
        return """<T-Digest: n=%d, centroids=%d>""" % (self.n, len(self))

    def __iter__(self):
        """
        Iterates over centroids in the digest.
        """
        return iter(self.centroids_to_list())

    def _add_centroid(self, centroid):
        if centroid.mean not in self.C:
            self.C.insert(centroid.mean, centroid)
        else:
            self.C[centroid.mean].update(centroid.mean, centroid.count)

    def _compute_centroid_quantile(self, centroid):
        denom = self.n
        cumulative_sum = self.C.get_left_accumulation(centroid.mean)
        return (centroid.count / 2. + cumulative_sum) / denom

    def _update_centroid(self, centroid, x, w):
        self.C.pop(centroid.mean)
        centroid.update(x, w)
        self._add_centroid(centroid)

    def _find_closest_centroids(self, x):
        try:
            ceil_key = self.C.ceiling_key(x)
        except KeyError:
            floor_key = self.C.floor_key(x)
            return [self.C[floor_key]]

        try:
            floor_key = self.C.floor_key(x)
        except KeyError:
            ceil_key = self.C.ceiling_key(x)
            return [self.C[ceil_key]]

        if abs(floor_key - x) < abs(ceil_key - x):
            return [self.C[floor_key]]
        elif abs(floor_key - x) == abs(ceil_key - x) and (ceil_key != floor_key):
            return [self.C[ceil_key], self.C[floor_key]]
        else:
            return [self.C[ceil_key]]

    def _threshold(self, q):
        return 4 * self.n * self.delta * q * (1 - q)

    def update(self, x, w=1):
        """
        Update the t-digest with value x and weight w.

        """
        self.n += w

        if len(self) == 0:
            self._add_centroid(Centroid(x, w))
            return

        S = self._find_closest_centroids(x)

        while len(S) != 0 and w > 0:
            j = choice(list(range(len(S))))
            c_j = S[j]

            q = self._compute_centroid_quantile(c_j)

            # This filters the out centroids that do not satisfy the second part
            # of the definition of S. See original paper by Dunning.
            if c_j.count + w > self._threshold(q):
                S.pop(j)
                continue

            delta_w = min(self._threshold(q) - c_j.count, w)
            self._update_centroid(c_j, x, delta_w)
            w -= delta_w
            S.pop(j)

        if w > 0:
            self._add_centroid(Centroid(x, w))

        if len(self) > self.K / self.delta:
            self.compress()

        return

    def batch_update(self, values, w=1):
        """
        Update the t-digest with an iterable of values. This assumes all points have the
        same weight.
        """
        for x in values:
            self.update(x, w)
        self.compress()
        return

    def compress(self):
        T = TDigest(self.delta, self.K)
        C = list(self.C.values())
        for c_i in pyudorandom.items(C):
            T.update(c_i.mean, c_i.count)
        self.C = T.C

    def percentile(self, p):
        """
        Computes the percentile of a specific value in [0,100].
        """

        if not (0 <= p <= 100):
            raise ValueError("p must be between 0 and 100, inclusive.")

        p = float(p)/100.
        p *= self.n
        c_i = None
        t = 0

        if p == 0:
            return self.C.min_item()[1].mean

        for i, key in enumerate(self.C.keys()):
            c_i_plus_one = self.C[key]
            if i == 0:
                k = c_i_plus_one.count / 2

            else:
                k = (c_i_plus_one.count + c_i.count) / 2.
                if p <= t + k:
                    z1 = p - t
                    z2 = t + k - p
                    return (c_i.mean * z2 + c_i_plus_one.mean * z1) / k
            c_i = c_i_plus_one
            t += k

        return self.C.max_item()[1].mean

    def cdf(self, x):
        """
        Computes the cdf of a specific value, ie. computes F(x) where F denotes
        the CDF of the distribution.
        """
        t = 0
        N = float(self.n)

        if len(self) == 1:  # only one centroid
            return int(x >= self.C.min_key())

        for i, key in enumerate(self.C.keys()):
            c_i = self.C[key]
            if i == len(self) - 1:
                delta = (c_i.mean - self.C.prev_item(key)[1].mean) / 2.
            else:
                delta = (self.C.succ_item(key)[1].mean - c_i.mean) / 2.
            z = max(-1, (x - c_i.mean) / delta)

            if z < 1:
                return t / N + c_i.count / N * (z + 1) / 2

            t += c_i.count
        return 1

    def trimmed_mean(self, p1, p2):
        """
        Computes the mean of the distribution between the two percentiles p1 and p2.
        This is a modified algorithm than the one presented in the original t-Digest paper.

        """
        if not (p1 < p2):
            raise ValueError("p1 must be between 0 and 100 and less than p2.")

        min_count = p1 / 100. * self.n
        max_count = p2 / 100. * self.n

        trimmed_sum = trimmed_count = curr_count = 0
        for i, c in enumerate(self.C.values()):
            next_count = curr_count + c.count
            if next_count <= min_count:
                curr_count = next_count
                continue

            count = c.count
            if curr_count < min_count:
                count = next_count - min_count
            if next_count > max_count:
                count -= next_count - max_count

            trimmed_sum += count * c.mean
            trimmed_count += count

            if next_count >= max_count:
                break
            curr_count = next_count

        if trimmed_count == 0:
            return 0
        return trimmed_sum / trimmed_count

    def centroids_to_list(self):
        """
        Returns a Python list of the TDigest object's Centroid values.

        """
        centroids = []
        for key in self.C.keys():
            tree_values = self.C.get_value(key)
            centroids.append({'m':tree_values.mean, 'c':tree_values.count})
        return centroids

    def to_dict(self):
        """
        Returns a Python dictionary of the TDigest and internal Centroid values.
        Or use centroids_to_list() for a list of only the Centroid values.

        """
        return {'n':self.n, 'delta':self.delta, 'K':self.K, 'centroids':self.centroids_to_list()}

    def update_from_dict(self, dict_values):
        """
        Updates TDigest object with dictionary values.

        The digest delta and K values are optional if you would like to update them,
        but the n value is not required because it is computed from the centroid weights.

        For example, you can initalize a new TDigest:
            digest = TDigest()
        Then load dictionary values into the digest:
            digest.update_from_dict({'K': 25, 'delta': 0.01, 'centroids': [{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}]})

        Or update an existing digest where the centroids will be appropriately merged:
            digest = TDigest()
            digest.update(1)
            digest.update(2)
            digest.update(3)
            digest.update_from_dict({'K': 25, 'delta': 0.01, 'centroids': [{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}]})

        Resulting in the digest having merged similar centroids by increasing their weight:
            {'K': 25, 'delta': 0.01, 'centroids': [{'c': 2.0, 'm': 1.0}, {'c': 2.0, 'm': 2.0}, {'c': 2.0, 'm': 3.0}], 'n': 6.0}

        Alternative you can provide only a list of centroid values with update_centroids_from_list()

        """
        self.delta = dict_values.get('delta', self.delta)
        self.K = dict_values.get('K', self.K)
        self.update_centroids_from_list(dict_values['centroids'])
        return self

    def update_centroids_from_list(self, list_values):
        """
        Add or update Centroids from a Python list.
        Any existing centroids in the digest object are appropriately updated.

        Example:
            digest.update_centroids([{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}])

        """
        [self.update(value['m'], value['c']) for value in list_values]
        return self

if __name__ == '__main__':
    from numpy import random

    T1 = TDigest()
    x = random.random(size=10000)
    T1.batch_update(x)

    print(abs(T1.percentile(50) - 0.5))
    print(abs(T1.percentile(10) - .1))
    print(abs(T1.percentile(90) - 0.9))
    print(abs(T1.percentile(1) - 0.01))
    print(abs(T1.percentile(0.1) - 0.001))
    print(T1.trimmed_mean(0.5, 1.))
<commit_msg>Update tdigest.py<commit_after>from __future__ import print_function

from random import choice
from itertools import chain
from accumulation_tree import AccumulationTree
import pyudorandom

def _centroid_count(centroid):
    # this is defined at the top level because
    # pyspark needs to pickle these functions properly
    # See https://github.com/CamDavidsonPilon/tdigest/issues/42
    return centroid.count

class Centroid(object):

    def __init__(self, mean, count):
        self.mean = float(mean)
        self.count = float(count)

    def __repr__(self):
        return """<Centroid: mean=%.8f, count=%d>""" % (self.mean, self.count)

    def __eq__(self, other):
        return self.mean == other.mean and self.count == other.count

    def update(self, x, weight):
        self.count += weight
        self.mean += weight * (x - self.mean) / self.count
        return

class TDigest(object):

    def __init__(self, delta=0.01, K=25):

        self.C = AccumulationTree(_centroid_count)
        self.n = 0
        self.delta = delta
        self.K = K

    def __add__(self, other_digest):
        data = list(chain(self.C.values(), other_digest.C.values()))
        new_digest = TDigest(self.delta, self.K)

        if len(data) > 0:
            for c in pyudorandom.items(data):
                new_digest.update(c.mean, c.count)

        return new_digest

    def __len__(self):
        return len(self.C)

    def __repr__(self):
        return """<T-Digest: n=%d, centroids=%d>""" % (self.n, len(self))

    def __iter__(self):
        """
        Iterates over centroids in the digest.
        """
        return iter(self.centroids_to_list())

    def _add_centroid(self, centroid):
        if centroid.mean not in self.C:
            self.C.insert(centroid.mean, centroid)
        else:
            self.C[centroid.mean].update(centroid.mean, centroid.count)

    def _compute_centroid_quantile(self, centroid):
        denom = self.n
        cumulative_sum = self.C.get_left_accumulation(centroid.mean)
        return (centroid.count / 2. + cumulative_sum) / denom

    def _update_centroid(self, centroid, x, w):
        self.C.pop(centroid.mean)
        centroid.update(x, w)
        self._add_centroid(centroid)

    def _find_closest_centroids(self, x):
        try:
            ceil_key = self.C.ceiling_key(x)
        except KeyError:
            floor_key = self.C.floor_key(x)
            return [self.C[floor_key]]

        try:
            floor_key = self.C.floor_key(x)
        except KeyError:
            ceil_key = self.C.ceiling_key(x)
            return [self.C[ceil_key]]

        if abs(floor_key - x) < abs(ceil_key - x):
            return [self.C[floor_key]]
        elif abs(floor_key - x) == abs(ceil_key - x) and (ceil_key != floor_key):
            return [self.C[ceil_key], self.C[floor_key]]
        else:
            return [self.C[ceil_key]]

    def _threshold(self, q):
        return 4 * self.n * self.delta * q * (1 - q)

    def update(self, x, w=1):
        """
        Update the t-digest with value x and weight w.

        """
        self.n += w

        if len(self) == 0:
            self._add_centroid(Centroid(x, w))
            return

        S = self._find_closest_centroids(x)

        while len(S) != 0 and w > 0:
            j = choice(list(range(len(S))))
            c_j = S[j]

            q = self._compute_centroid_quantile(c_j)

            # This filters the out centroids that do not satisfy the second part
            # of the definition of S. See original paper by Dunning.
            if c_j.count + w > self._threshold(q):
                S.pop(j)
                continue

            delta_w = min(self._threshold(q) - c_j.count, w)
            self._update_centroid(c_j, x, delta_w)
            w -= delta_w
            S.pop(j)

        if w > 0:
            self._add_centroid(Centroid(x, w))

        if len(self) > self.K / self.delta:
            self.compress()

        return

    def batch_update(self, values, w=1):
        """
        Update the t-digest with an iterable of values. This assumes all points have the
        same weight.
        """
        for x in values:
            self.update(x, w)
        self.compress()
        return

    def compress(self):
        T = TDigest(self.delta, self.K)
        C = list(self.C.values())
        for c_i in pyudorandom.items(C):
            T.update(c_i.mean, c_i.count)
        self.C = T.C

    def percentile(self, p):
        """
        Computes the percentile of a specific value in [0,100].
        """

        if not (0 <= p <= 100):
            raise ValueError("p must be between 0 and 100, inclusive.")

        p = float(p)/100.
        p *= self.n
        c_i = None
        t = 0

        if p == 0:
            return self.C.min_item()[1].mean

        for i, key in enumerate(self.C.keys()):
            c_i_plus_one = self.C[key]
            if i == 0:
                k = c_i_plus_one.count / 2

            else:
                k = (c_i_plus_one.count + c_i.count) / 2.
                if p <= t + k:
                    z1 = p - t
                    z2 = t + k - p
                    return (c_i.mean * z2 + c_i_plus_one.mean * z1) / k
            c_i = c_i_plus_one
            t += k

        return self.C.max_item()[1].mean

    def cdf(self, x):
        """
        Computes the cdf of a specific value, ie. computes F(x) where F denotes
        the CDF of the distribution.
        """
        t = 0
        N = float(self.n)

        if len(self) == 1:  # only one centroid
            return int(x >= self.C.min_key())

        for i, key in enumerate(self.C.keys()):
            c_i = self.C[key]
            if i == len(self) - 1:
                delta = (c_i.mean - self.C.prev_item(key)[1].mean) / 2.
            else:
                delta = (self.C.succ_item(key)[1].mean - c_i.mean) / 2.
            z = max(-1, (x - c_i.mean) / delta)

            if z < 1:
                return t / N + c_i.count / N * (z + 1) / 2

            t += c_i.count
        return 1

    def trimmed_mean(self, p1, p2):
        """
        Computes the mean of the distribution between the two percentiles p1 and p2.
        This is a modified algorithm than the one presented in the original t-Digest paper.

        """
        if not (p1 < p2):
            raise ValueError("p1 must be between 0 and 100 and less than p2.")

        min_count = p1 / 100. * self.n
        max_count = p2 / 100. * self.n

        trimmed_sum = trimmed_count = curr_count = 0
        for i, c in enumerate(self.C.values()):
            next_count = curr_count + c.count
            if next_count <= min_count:
                curr_count = next_count
                continue

            count = c.count
            if curr_count < min_count:
                count = next_count - min_count
            if next_count > max_count:
                count -= next_count - max_count

            trimmed_sum += count * c.mean
            trimmed_count += count

            if next_count >= max_count:
                break
            curr_count = next_count

        if trimmed_count == 0:
            return 0
        return trimmed_sum / trimmed_count

    def centroids_to_list(self):
        """
        Returns a Python list of the TDigest object's Centroid values.

        """
        centroids = []
        for key in self.C.keys():
            tree_values = self.C.get_value(key)
            centroids.append({'m':tree_values.mean, 'c':tree_values.count})
        return centroids

    def to_dict(self):
        """
        Returns a Python dictionary of the TDigest and internal Centroid values.
        Or use centroids_to_list() for a list of only the Centroid values.

        """
        return {'n':self.n, 'delta':self.delta, 'K':self.K, 'centroids':self.centroids_to_list()}

    @staticmethod
    def from_dict(d:dict):
       t = TDigest()
       for c in d['centroids']:
           ci = Centroid(c['m'],c['c'])
           t.C.insert(c['m'], ci)
       t.K = d['K']
       t.delta = d['delta']
       t.n = d['n']
       return t

    def update_from_dict(self, dict_values):
        """
        Updates TDigest object with dictionary values.

        The digest delta and K values are optional if you would like to update them,
        but the n value is not required because it is computed from the centroid weights.

        For example, you can initalize a new TDigest:
            digest = TDigest()
        Then load dictionary values into the digest:
            digest.update_from_dict({'K': 25, 'delta': 0.01, 'centroids': [{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}]})

        Or update an existing digest where the centroids will be appropriately merged:
            digest = TDigest()
            digest.update(1)
            digest.update(2)
            digest.update(3)
            digest.update_from_dict({'K': 25, 'delta': 0.01, 'centroids': [{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}]})

        Resulting in the digest having merged similar centroids by increasing their weight:
            {'K': 25, 'delta': 0.01, 'centroids': [{'c': 2.0, 'm': 1.0}, {'c': 2.0, 'm': 2.0}, {'c': 2.0, 'm': 3.0}], 'n': 6.0}

        Alternative you can provide only a list of centroid values with update_centroids_from_list()

        """
        self.delta = dict_values.get('delta', self.delta)
        self.K = dict_values.get('K', self.K)
        self.update_centroids_from_list(dict_values['centroids'])
        return self

    def update_centroids_from_list(self, list_values):
        """
        Add or update Centroids from a Python list.
        Any existing centroids in the digest object are appropriately updated.

        Example:
            digest.update_centroids([{'c': 1.0, 'm': 1.0}, {'c': 1.0, 'm': 2.0}, {'c': 1.0, 'm': 3.0}])

        """
        [self.update(value['m'], value['c']) for value in list_values]
        return self

if __name__ == '__main__':
    from numpy import random

    T1 = TDigest()
    x = random.random(size=10000)
    T1.batch_update(x)

    print(abs(T1.percentile(50) - 0.5))
    print(abs(T1.percentile(10) - .1))
    print(abs(T1.percentile(90) - 0.9))
    print(abs(T1.percentile(1) - 0.01))
    print(abs(T1.percentile(0.1) - 0.001))
    print(T1.trimmed_mean(0.5, 1.))
<|endoftext|>


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

'''
This script does two things:
1) Ouputs strings that haven't been translated in files for each language
2) Cleans up string.xml files for other languages by removing old strings and
re-ordering the strings based on their order in the default language

Other important notes
-Should support all language codes that follow the -** or -**-*** pattern
-This ignores strings that start with "provider." or have the "translatable"
attribute set to "false"
-The output directory for the missing strings is in the current directory
-The output file names look like "strings_to_trans-**" where "**" is the language
code
'''

import sys
import os
import xml.etree.ElementTree as ET
import codecs
import argparse

ORIG_DIR = os.getcwd()

def main():

    # parse command line arguments
    res_path, clean, out_path = parseArgs()

    # go to the resource directory and save the whole path
    os.chdir(res_path)
    res_path = os.getcwd()

    # get default keys
    tree = getDefaultTree(res_path)
    keys = getKeysFromTree(tree)
    tags = getTagsFromTree(tree)

    # get the languages that we want to translate to
    langs = getLangsFromDir(res_path)

    # look for missing keys in each language string file
    missing = findMissingKeys(keys, langs, res_path)

    # remove old strings and sort them all in the same way
    if (clean):
        cleanTranslationFiles(langs, keys, res_path)

    # write files for missing keys for each language
    createOutputDir(out_path)
    writeMissingKeysToFiles(langs, tags, missing, out_path)

def parseArgs():
    # parse arguments and do error checking
    parser = argparse.ArgumentParser()
    parser.add_argument('--res',
                        help='Path to the app\'s /res directory. If not specifies it assumes current directory',
                        default='.')
    parser.add_argument('--output',
                        help='Path to the output directory. If not specifies it will create a folder called to_translate in the current directory',
                        default='./to_translate')
    parser.add_argument("--clean", help="re-orders and removes strings in the translation files to match the default string ordering",
                    action="store_true")
    args = parser.parse_args()
    return args.res, args.clean, args.output

def getDefaultTree(res_path):
    os.chdir(res_path)
    os.chdir('values')
    ET.register_namespace('tools', "http://schemas.android.com/tools")
    ET.register_namespace('xliff', "urn:oasis:names:tc:xliff:document:1.2")
    return ET.parse('strings.xml')

def createOutputDir(out_path):
    # create output directory
    os.chdir(ORIG_DIR)
    if not os.path.exists(out_path):
        os.makedirs(out_path)

def writeMissingKeysToFiles(langs, tags, missing, out_path):
    # write xml files for missing strings for each language
    os.chdir(ORIG_DIR)
    os.chdir(out_path)
    for lang in langs:
        # skip language if it's not missing any strings
        if (len(missing[lang]) == 0):
            continue

        # create element tree for all the missing tags
        root = ET.Element('resources')
        for key in missing[lang]:
            tag = getTagByKeyName(tags, key)
            root.append(tag)

        # write out the strings
        f = codecs.open('strings_to_trans-%s.xml' % (lang), 'wb', 'utf-8')
        f.write(prettify(root))

def getLanguageTrees(langs, res_path):
    trees = {}
    for lang in langs:
        os.chdir(res_path)
        os.chdir('values-' + lang)
        trees[lang] = ET.parse('strings.xml')
    return trees

def cleanTranslationFiles(langs, keys, res_path):
    trees = getLanguageTrees(langs, res_path)
    for lang in trees.keys():
        tree = trees[lang]
        keys_trans = getKeysFromTree(tree)
        tags_trans = getTagsFromTree(tree)
        keys_has = []
        for key in keys:
            if key in keys_trans:
                keys_has.append(key)
        root = ET.Element('resources')
        for key in keys_has:
            tag = getTagByKeyName(tags_trans, key)
            root.append(tag)

        # write out file
        os.chdir(res_path)
        os.chdir('values-%s' % (lang))
        f = codecs.open('strings.xml', 'wb', 'utf-8')
        print(prettify(root))
        f.write(prettify(root))

def getTagByKeyName(tags, key):
    for tag in tags:
        if tag.get('name') == key:
            return tag

'''
Return a pretty-printed XML string for the Element.
'''
def prettify(elem):
    TAB = '    '
    elems = ET.tostringlist(elem, encoding='UTF-8')
    output = ''
    for i in range(len(elems)):
        elem = elems[i].decode('utf-8')
        if i > 0:
            prev1 = elems[i-1].decode('utf-8')
        if i > 1:
            prev2 = elems[i-2].decode('utf-8')
        # make sure strings and plurals are indented properly
        # (everything else should be fine since it is nested within those tags)
        if elem == '<string' or elem == '<plurals':
            if not prev2.startswith('\n') and not prev1.startswith('\n'):
                output += '\n'
            if not prev1.endswith(TAB):
                output += TAB

        output += elem
    return output

def findMissingKeys(keys, langs, res_path):
    missing = {}
    trees = getLanguageTrees(langs, res_path)
    for lang in trees.keys():
        tree = trees[lang]
        keys_trans = getKeysFromTree(tree)
        keys_miss = []
        for key in keys:
            if key not in keys_trans:
                keys_miss.append(key)
        missing[lang] = keys_miss
    return missing

'''
Supported langauge directories follow one of two patterns:
1) values-**
2) values-**-***

returns code for language or None if not a language directory
'''
def getLangDir(dir_name):
    if dir_name[2:].startswith('values-'):
        code = [dir_name[9:]][0]
        if (len(code) == 2) or (len(code) == 6 and code[2] == '-'):
            return code

    # not a language dir
    return None

def getLangsFromDir(res_path):
    os.chdir(res_path)
    langs = []
    for x in os.walk('.'):
        code = getLangDir(x[0])
        if code is not None:
            langs.append(code)
    return langs

def getKeysFromTree(tree):
    root = tree.getroot()
    keys = []
    for child in root:
        # ignore strings that can't be translated
        if child.get('translatable', default='true') == 'false':
            continue
        # ignore providers
        if (child.get('name').startswith('provider.')):
            continue
        keys.append(child.get('name'))
    return keys

def getTagsFromTree(tree):
    root = tree.getroot()
    keys = []
    for child in root:
        keys.append(child)
    return keys

if __name__ == '__main__':
    main()

<commit_msg>barebones fix<commit_after>#!/usr/bin/env python

'''
This script does two things:
1) Ouputs strings that haven't been translated in files for each language
2) Cleans up string.xml files for other languages by removing old strings and
re-ordering the strings based on their order in the default language

Other important notes
-Should support all language codes that follow the -** or -**-*** pattern
-This ignores strings that start with "provider." or have the "translatable"
attribute set to "false"
-The output directory for the missing strings is in the current directory
-The output file names look like "strings_to_trans-**" where "**" is the language
code
'''

import sys
import os
import xml.etree.ElementTree as ET
import codecs
import argparse

ORIG_DIR = os.getcwd()

def main():

    # parse command line arguments
    res_path, clean, out_path = parseArgs()

    # go to the resource directory and save the whole path
    os.chdir(res_path)
    res_path = os.getcwd()

    # get default keys
    tree = getDefaultTree(res_path)
    keys = getKeysFromTree(tree)
    tags = getTagsFromTree(tree)

    # get the languages that we want to translate to
    langs = getLangsFromDir(res_path)

    # look for missing keys in each language string file
    missing = findMissingKeys(keys, langs, res_path)

    # remove old strings and sort them all in the same way
    if (clean):
        cleanTranslationFiles(langs, keys, res_path)

    # write files for missing keys for each language
    createOutputDir(out_path)
    writeMissingKeysToFiles(langs, tags, missing, out_path)

def parseArgs():
    # parse arguments and do error checking
    parser = argparse.ArgumentParser()
    parser.add_argument('--res',
                        help='Path to the app\'s /res directory. If not specifies it assumes current directory',
                        default='.')
    parser.add_argument('--output',
                        help='Path to the output directory. If not specifies it will create a folder called to_translate in the current directory',
                        default='./to_translate')
    parser.add_argument("--clean", help="re-orders and removes strings in the translation files to match the default string ordering",
                    action="store_true")
    args = parser.parse_args()
    return args.res, args.clean, args.output

def getDefaultTree(res_path):
    os.chdir(res_path)
    os.chdir('values')
    ET.register_namespace('tools', "http://schemas.android.com/tools")
    ET.register_namespace('xliff', "urn:oasis:names:tc:xliff:document:1.2")
    return ET.parse('strings.xml')

def createOutputDir(out_path):
    # create output directory
    os.chdir(ORIG_DIR)
    if not os.path.exists(out_path):
        os.makedirs(out_path)

def writeMissingKeysToFiles(langs, tags, missing, out_path):
    # write xml files for missing strings for each language
    os.chdir(ORIG_DIR)
    os.chdir(out_path)
    for lang in langs:
        # skip language if it's not missing any strings
        if (len(missing[lang]) == 0):
            continue

        # create element tree for all the missing tags
        root = ET.Element('resources')
        for key in missing[lang]:
            tag = getTagByKeyName(tags, key)
            root.append(tag)

        # write out the strings
        f = codecs.open('strings_to_trans-%s.xml' % (lang), 'wb', 'utf-8')
        f.write(prettify(root))

def getLanguageTrees(langs, res_path):
    trees = {}
    for lang in langs:
        os.chdir(res_path)
        os.chdir('values-' + lang)
        trees[lang] = ET.parse('strings.xml')
    return trees

def cleanTranslationFiles(langs, keys, res_path):
    trees = getLanguageTrees(langs, res_path)
    for lang in trees.keys():
        tree = trees[lang]
        keys_trans = getKeysFromTree(tree)
        tags_trans = getTagsFromTree(tree)
        keys_has = []
        for key in keys:
            if key in keys_trans:
                keys_has.append(key)
        root = ET.Element('resources')
        for key in keys_has:
            tag = getTagByKeyName(tags_trans, key)
            root.append(tag)

        # write out file
        os.chdir(res_path)
        os.chdir('values-%s' % (lang))
        f = codecs.open('strings.xml', 'wb', 'utf-8')
        f.write(prettify(root))

def getTagByKeyName(tags, key):
    for tag in tags:
        if tag.get('name') == key:
            return tag

'''
Return a pretty-printed XML string for the Element.
'''
def prettify(elem):
    output = ET.tostring(elem, encoding='UTF-8').decode('utf-8')

    # make sure we add the xml declaration... stupid python 3
    if not output.startswith('<?xml'):
        output = "<?xml version='1.0' encoding='UTF-8'?>\n" + output

    # fix first string not indenting
    output = output.replace('><string', '>\n    <string')
    return output

def findMissingKeys(keys, langs, res_path):
    missing = {}
    trees = getLanguageTrees(langs, res_path)
    for lang in trees.keys():
        tree = trees[lang]
        keys_trans = getKeysFromTree(tree)
        keys_miss = []
        for key in keys:
            if key not in keys_trans:
                keys_miss.append(key)
        missing[lang] = keys_miss
    return missing

'''
Supported langauge directories follow one of two patterns:
1) values-**
2) values-**-***

returns code for language or None if not a language directory
'''
def getLangDir(dir_name):
    if dir_name[2:].startswith('values-'):
        code = [dir_name[9:]][0]
        if (len(code) == 2) or (len(code) == 6 and code[2] == '-'):
            return code

    # not a language dir
    return None

def getLangsFromDir(res_path):
    os.chdir(res_path)
    langs = []
    for x in os.walk('.'):
        code = getLangDir(x[0])
        if code is not None:
            langs.append(code)
    return langs

def getKeysFromTree(tree):
    root = tree.getroot()
    keys = []
    for child in root:
        # ignore strings that can't be translated
        if child.get('translatable', default='true') == 'false':
            continue
        # ignore providers
        if (child.get('name').startswith('provider.')):
            continue
        keys.append(child.get('name'))
    return keys

def getTagsFromTree(tree):
    root = tree.getroot()
    keys = []
    for child in root:
        keys.append(child)
    return keys

if __name__ == '__main__':
    main()

<|endoftext|>