<commit_before># Copyright 2011, 2012 Omniscale (http://omniscale.com)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division

import math
import codecs
import os
import shapely.geometry
import shapely.geos
import shapely.prepared
from shapely.geometry.base import BaseGeometry
from shapely import geometry
from shapely import wkt
from shapely.ops import cascaded_union, linemerge
from shapely.topology import TopologicalError

try:
    import rtree
except ImportError:
    rtree = None

from imposm import config
from imposm.util.geom import load_polygons, load_datasource, build_multipolygon

import logging
log = logging.getLogger(__name__)

class InvalidGeometryError(Exception):
    pass

class IncompletePolygonError(Exception):
    pass

TOLERANCE_DEEGREES = 1e-8
TOLERANCE_METERS = 1e-3

# older versions had unhandled floating point execptions in .buffer(0)
SHAPELY_SUPPORTS_BUFFER = shapely.geos.geos_capi_version >= (1, 6, 0)

def validate_and_simplify(geom, meter_units=False):
    if SHAPELY_SUPPORTS_BUFFER:
        try:
            # buffer(0) is nearly fast as is_valid
            return geom.buffer(0)
        except ValueError:
            # shapely raises ValueError if buffer(0) result is empty
            raise InvalidGeometryError('geometry is empty')

    orig_geom = geom
    if not geom.is_valid:
        tolerance = TOLERANCE_METERS if meter_units else TOLERANCE_DEEGREES
        try:
            geom = geom.simplify(tolerance, False)
        except ValueError:
            # shapely raises ValueError if buffer(0) result is empty
            raise InvalidGeometryError('geometry is empty')

        if not geom.is_valid:
            raise InvalidGeometryError('geometry is invalid, could not simplify: %s' %
                                       orig_geom)
    return geom

class GeomBuilder(object):
    def build(self, osm_elem):
        # TODO is this method still in use?
        try:
            if isinstance(osm_elem.coords, BaseGeometry):
                return osm_elem.coords
            geom_wkt = self.to_wkt(osm_elem.coords)
            if geom_wkt is not None:
                geom = wkt.loads(geom_wkt)
        except Exception, ex:
            raise InvalidGeometryError('unable to build geometry %s: %s %s' %
                                       (osm_elem.osm_id, ex, osm_elem.coords))
        if geom_wkt is None or geom is None:
            # unable to build valid wkt (non closed polygon, etc)
            raise InvalidGeometryError()
        return geom

    def check_geom_type(self, geom):
        return

    def build_geom(self, osm_elem):
        try:
            if isinstance(osm_elem.coords, BaseGeometry):
                if osm_elem.coords.is_empty:
                    raise InvalidGeometryError('empty geometry')
                self.check_geom_type(osm_elem.coords)
                return osm_elem.coords
            geom = self.to_geom(osm_elem.coords)
        except InvalidGeometryError, ex:
            raise InvalidGeometryError('unable to build geometry %s: %s %s' %
                                       (osm_elem.osm_id, ex, osm_elem.coords))
        if geom is None:
            # unable to build valid wkt (non closed polygon, etc)
            raise InvalidGeometryError()
        return geom

class PointBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) != 2:
            return None
        return 'POINT(%f %f)' % data

    def to_geom(self, data):
        if len(data) != 2:
            return None
        return geometry.Point(*data)

    def check_geom_type(self, geom):
        if geom.type != 'Point':
            raise InvalidGeometryError('expected Point, got %s' % geom.type)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate or geom.is_valid:
            return geom
        else:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

class PolygonBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) >= 4 and data[0] == data[-1]:
            return 'POLYGON((' + ', '.join('%f %f' % p for p in data) + '))'
        return None

    def to_geom(self, data):
        if len(data) >= 4 and data[0] == data[-1]:
            return geometry.Polygon(data)
        return None

    def check_geom_type(self, geom):
        if geom.type not in ('Polygon', 'MultiPolygon'):
            raise InvalidGeometryError('expected Polygon or MultiPolygon, got %s' % geom.type)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate:
            return geom
        try:
            return validate_and_simplify(geom)
        except InvalidGeometryError:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

class LineStringBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) <= 1:
            return None
        if len(data) == 2 and data[0] == data[1]:
            return None
        return 'LINESTRING(' + ', '.join('%f %f' % p for p in data) + ')'

    def check_geom_type(self, geom):
        if geom.type != 'LineString':
            raise InvalidGeometryError('expected LineString, got %s' % geom.type)

    def to_geom(self, data, max_length=None):
        if len(data) <= 1:
            return None
        if len(data) == 2 and data[0] == data[1]:
            return None
        if max_length is None:
            max_length = config.imposm_linestring_max_length
        if max_length and len(data) > max_length:
            chunks = math.ceil(len(data) / max_length)
            length = int(len(data) // chunks)
            lines = []
            for i in xrange(1, len(data), length):
                lines.append(geometry.LineString(data[i-1:i+length]))
            return lines
        return geometry.LineString(data)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate or geom.is_valid:
            return geom
        else:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

def tile_bbox(bbox, grid_width):
    """
    Tile bbox into multiple sub-boxes, each of `grid_width` size.

    >>> list(tile_bbox((-1, 1, 0.49, 1.51), 0.5)) #doctest: +NORMALIZE_WHITESPACE
    [(-1.0, 1.0, -0.5, 1.5),
     (-1.0, 1.5, -0.5, 2.0),
     (-0.5, 1.0, 0.0, 1.5),
     (-0.5, 1.5, 0.0, 2.0),
     (0.0, 1.0, 0.5, 1.5),
     (0.0, 1.5, 0.5, 2.0)]
    """
    min_x = math.floor(bbox[0]/grid_width) * grid_width
    min_y = math.floor(bbox[1]/grid_width) * grid_width
    max_x = math.ceil(bbox[2]/grid_width) * grid_width
    max_y = math.ceil(bbox[3]/grid_width) * grid_width

    x_steps = (max_x - min_x) / grid_width
    y_steps = (max_y - min_y) / grid_width

    for x in xrange(int(x_steps)):
        for y in xrange(int(y_steps)):
            yield (
                min_x + x * grid_width,
                min_y + y * grid_width,
                min_x + (x + 1) * grid_width,
                min_y + (y + 1)* grid_width,
            )

def split_polygon_at_grid(geom, grid_width=0.1):
    """
    >>> p = list(split_polygon_at_grid(geometry.box(-0.5, 1, 0.2, 2), 1))
    >>> p[0].contains(geometry.box(-0.5, 1, 0, 2))
    True
    >>> p[0].area == geometry.box(-0.5, 1, 0, 2).area
    True
    >>> p[1].contains(geometry.box(0, 1, 0.2, 2))
    True
    >>> p[1].area == geometry.box(0, 1, 0.2, 2).area
    True
    """
    if not geom.is_valid:
        geom = geom.buffer(0)
    for split_box in tile_bbox(geom.bounds, grid_width):
        try:
            polygon_part = geom.intersection(shapely.geometry.box(*split_box))
        except TopologicalError:
            continue
        if not polygon_part.is_empty:
            yield polygon_part

def load_geom(source):
    geom = None
    # source is a wkt file
    if os.path.exists(os.path.abspath(source)):
        data = None
        with open(os.path.abspath(source), 'r') as fp:
            data = fp.read(50)
        # load WKT geometry and remove leading whitespaces
        if data.lower().lstrip().startswith(('polygon', 'multipolygon')):
            geom = load_polygons(source)
    # source is an OGR datasource
    if geom is None:
        geom = load_datasource(source)

    if geom:
        if rtree:
            return LimitRTreeGeometry(geom)
        else:
            log.info('You should install RTree for large --limit-to polygons')
            return LimitPolygonGeometry(build_multipolygon(geom)[1])
    return None

class EmtpyGeometryError(Exception):
    pass

class LimitPolygonGeometry(object):
    def __init__(self, shapely_geom):
        self._geom = shapely_geom
        self._prepared_geom = None
        self._prepared_counter = 0
        self._prepared_max = 100000

    @property
    def geom(self):
        # GEOS internal data structure for prepared geometries grows over time,
        # recreate to limit memory consumption
        if not self._prepared_geom or self._prepared_counter > self._prepared_max:
            print 'create prepared'
            self._prepared_geom = shapely.prepared.prep(self._geom)
            self._prepared_counter = 0
        self._prepared_counter += 1
        return self._prepared_geom

    def intersection(self, geom):
        if self.geom.contains_properly(geom):
            # no need to limit contained geometries
            return geom

        new_geom = None
        if self.geom.intersects(geom):
            try:
                # can not use intersection with prepared geom
                new_geom = self._geom.intersection(geom)
            except TopologicalError:
                pass

        if not new_geom or new_geom.is_empty:
            raise EmtpyGeometryError('No intersection or empty geometry')

        new_geom = filter_geometry_by_type(new_geom, geom.type)
        if new_geom:
            return new_geom

        raise EmtpyGeometryError('No intersection or empty geometry')

def filter_geometry_by_type(geometry, geom_type):
    """
    Filter (multi)geometry for compatible `geom_type`,
    because we can't insert points into linestring tables for example
    """
    if geometry.type == geom_type:
        # same type is fine
        return geometry

    if geometry.type == 'MultiPolygon' and geom_type == 'Polygon':
        # polygon mappings should also support multipolygons
        return geometry

    if hasattr(geometry, 'geoms'):
        # GeometryCollection or MultiLineString? return list of geometries
        geoms = []
        for part in geometry.geoms:
            # only parts with same type
            if part.type == geom_type:
                geoms.append(part)

        if geoms:
            return geoms

    return None

def flatten_polygons(polygons):
    for polygon in polygons:
        if polygon.type == 'MultiPolygon':
            for p in polygon.geoms:
                yield p
        else:
            yield polygon

def flatten_linestrings(linestrings):
    for linestring in linestrings:
        if linestring.type == 'MultiLineString':
            for ls in linestring.geoms:
                yield ls
        else:
            yield linestring

class LimitRTreeGeometry(object):
    def __init__(self, polygons):
        index = rtree.index.Index()
        sub_polygons = []
        part_idx = 0
        for polygon in polygons:
            for part in split_polygon_at_grid(polygon):
                sub_polygons.append(part)
                index.insert(part_idx, part.bounds)
                part_idx += 1

        self.polygons = sub_polygons
        self.index = index

    def intersection(self, geom):
        intersection_ids = list(self.index.intersection(geom.bounds))

        if not intersection_ids:
            raise EmtpyGeometryError('No intersection or empty geometry')

        intersections = []
        for i in intersection_ids:

            polygon = self.polygons[i]

            if polygon.contains(geom):
                return geom

            if polygon.intersects(geom):
                try:
                    new_geom_part = polygon.intersection(geom)
                    new_geom_part = filter_geometry_by_type(new_geom_part, geom.type)
                    if new_geom_part:
                        if len(intersection_ids) == 1:
                            return new_geom_part

                        if isinstance(new_geom_part, list):
                            intersections.extend(new_geom_part)
                        else:
                            intersections.append(new_geom_part)
                except TopologicalError:
                        pass

        if not intersections:
            raise EmtpyGeometryError('No intersection or empty geometry')

        # intersections from multiple sub-polygons
        # try to merge them back to a single geometry
        if geom.type.endswith('Polygon'):
            union = cascaded_union(list(flatten_polygons(intersections)))
        elif geom.type.endswith('LineString'):
            union = linemerge(list(flatten_linestrings(intersections)))
            if union.type == 'MultiLineString':
                union = list(union.geoms)
        elif geom.type == 'Point':
            union = intersections[0]
        else:
            log.warn('unexpexted geometry type %s', geom.type)
            raise EmtpyGeometryError()
        return union
<commit_msg>added very simple srs check<commit_after># Copyright 2011, 2012 Omniscale (http://omniscale.com)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import division

import math
import codecs
import os
import shapely.geometry
import shapely.geos
import shapely.prepared
from shapely.geometry.base import BaseGeometry
from shapely import geometry
from shapely import wkt
from shapely.ops import cascaded_union, linemerge
from shapely.topology import TopologicalError

try:
    import rtree
except ImportError:
    rtree = None

from imposm import config
from imposm.util.geom import load_polygons, load_datasource, build_multipolygon

import logging
log = logging.getLogger(__name__)

class InvalidGeometryError(Exception):
    pass

class IncompletePolygonError(Exception):
    pass

TOLERANCE_DEEGREES = 1e-8
TOLERANCE_METERS = 1e-3

# older versions had unhandled floating point execptions in .buffer(0)
SHAPELY_SUPPORTS_BUFFER = shapely.geos.geos_capi_version >= (1, 6, 0)

def validate_and_simplify(geom, meter_units=False):
    if SHAPELY_SUPPORTS_BUFFER:
        try:
            # buffer(0) is nearly fast as is_valid
            return geom.buffer(0)
        except ValueError:
            # shapely raises ValueError if buffer(0) result is empty
            raise InvalidGeometryError('geometry is empty')

    orig_geom = geom
    if not geom.is_valid:
        tolerance = TOLERANCE_METERS if meter_units else TOLERANCE_DEEGREES
        try:
            geom = geom.simplify(tolerance, False)
        except ValueError:
            # shapely raises ValueError if buffer(0) result is empty
            raise InvalidGeometryError('geometry is empty')

        if not geom.is_valid:
            raise InvalidGeometryError('geometry is invalid, could not simplify: %s' %
                                       orig_geom)
    return geom

class GeomBuilder(object):
    def build(self, osm_elem):
        # TODO is this method still in use?
        try:
            if isinstance(osm_elem.coords, BaseGeometry):
                return osm_elem.coords
            geom_wkt = self.to_wkt(osm_elem.coords)
            if geom_wkt is not None:
                geom = wkt.loads(geom_wkt)
        except Exception, ex:
            raise InvalidGeometryError('unable to build geometry %s: %s %s' %
                                       (osm_elem.osm_id, ex, osm_elem.coords))
        if geom_wkt is None or geom is None:
            # unable to build valid wkt (non closed polygon, etc)
            raise InvalidGeometryError()
        return geom

    def check_geom_type(self, geom):
        return

    def build_geom(self, osm_elem):
        try:
            if isinstance(osm_elem.coords, BaseGeometry):
                if osm_elem.coords.is_empty:
                    raise InvalidGeometryError('empty geometry')
                self.check_geom_type(osm_elem.coords)
                return osm_elem.coords
            geom = self.to_geom(osm_elem.coords)
        except InvalidGeometryError, ex:
            raise InvalidGeometryError('unable to build geometry %s: %s %s' %
                                       (osm_elem.osm_id, ex, osm_elem.coords))
        if geom is None:
            # unable to build valid wkt (non closed polygon, etc)
            raise InvalidGeometryError()
        return geom

class PointBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) != 2:
            return None
        return 'POINT(%f %f)' % data

    def to_geom(self, data):
        if len(data) != 2:
            return None
        return geometry.Point(*data)

    def check_geom_type(self, geom):
        if geom.type != 'Point':
            raise InvalidGeometryError('expected Point, got %s' % geom.type)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate or geom.is_valid:
            return geom
        else:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

class PolygonBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) >= 4 and data[0] == data[-1]:
            return 'POLYGON((' + ', '.join('%f %f' % p for p in data) + '))'
        return None

    def to_geom(self, data):
        if len(data) >= 4 and data[0] == data[-1]:
            return geometry.Polygon(data)
        return None

    def check_geom_type(self, geom):
        if geom.type not in ('Polygon', 'MultiPolygon'):
            raise InvalidGeometryError('expected Polygon or MultiPolygon, got %s' % geom.type)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate:
            return geom
        try:
            return validate_and_simplify(geom)
        except InvalidGeometryError:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

class LineStringBuilder(GeomBuilder):
    def to_wkt(self, data):
        if len(data) <= 1:
            return None
        if len(data) == 2 and data[0] == data[1]:
            return None
        return 'LINESTRING(' + ', '.join('%f %f' % p for p in data) + ')'

    def check_geom_type(self, geom):
        if geom.type != 'LineString':
            raise InvalidGeometryError('expected LineString, got %s' % geom.type)

    def to_geom(self, data, max_length=None):
        if len(data) <= 1:
            return None
        if len(data) == 2 and data[0] == data[1]:
            return None
        if max_length is None:
            max_length = config.imposm_linestring_max_length
        if max_length and len(data) > max_length:
            chunks = math.ceil(len(data) / max_length)
            length = int(len(data) // chunks)
            lines = []
            for i in xrange(1, len(data), length):
                lines.append(geometry.LineString(data[i-1:i+length]))
            return lines
        return geometry.LineString(data)

    def build_checked_geom(self, osm_elem, validate=False):
        geom = self.build_geom(osm_elem)
        if not validate or geom.is_valid:
            return geom
        else:
            raise InvalidGeometryError('invalid geometry for %s: %s, %s' %
                                       (osm_elem.osm_id, geom, osm_elem.coords))

def tile_bbox(bbox, grid_width):
    """
    Tile bbox into multiple sub-boxes, each of `grid_width` size.

    >>> list(tile_bbox((-1, 1, 0.49, 1.51), 0.5)) #doctest: +NORMALIZE_WHITESPACE
    [(-1.0, 1.0, -0.5, 1.5),
     (-1.0, 1.5, -0.5, 2.0),
     (-0.5, 1.0, 0.0, 1.5),
     (-0.5, 1.5, 0.0, 2.0),
     (0.0, 1.0, 0.5, 1.5),
     (0.0, 1.5, 0.5, 2.0)]
    """
    min_x = math.floor(bbox[0]/grid_width) * grid_width
    min_y = math.floor(bbox[1]/grid_width) * grid_width
    max_x = math.ceil(bbox[2]/grid_width) * grid_width
    max_y = math.ceil(bbox[3]/grid_width) * grid_width

    x_steps = (max_x - min_x) / grid_width
    y_steps = (max_y - min_y) / grid_width

    for x in xrange(int(x_steps)):
        for y in xrange(int(y_steps)):
            yield (
                min_x + x * grid_width,
                min_y + y * grid_width,
                min_x + (x + 1) * grid_width,
                min_y + (y + 1)* grid_width,
            )

def split_polygon_at_grid(geom, grid_width=0.1):
    """
    >>> p = list(split_polygon_at_grid(geometry.box(-0.5, 1, 0.2, 2), 1))
    >>> p[0].contains(geometry.box(-0.5, 1, 0, 2))
    True
    >>> p[0].area == geometry.box(-0.5, 1, 0, 2).area
    True
    >>> p[1].contains(geometry.box(0, 1, 0.2, 2))
    True
    >>> p[1].area == geometry.box(0, 1, 0.2, 2).area
    True
    """
    if not geom.is_valid:
        geom = geom.buffer(0)
    for split_box in tile_bbox(geom.bounds, grid_width):
        try:
            polygon_part = geom.intersection(shapely.geometry.box(*split_box))
        except TopologicalError:
            continue
        if not polygon_part.is_empty:
            yield polygon_part

def load_geom(source):
    geom = None
    # source is a wkt file
    if os.path.exists(os.path.abspath(source)):
        data = None
        with open(os.path.abspath(source), 'r') as fp:
            data = fp.read(50)
        # load WKT geometry and remove leading whitespaces
        if data.lower().lstrip().startswith(('polygon', 'multipolygon')):
            geom = load_polygons(source)
    # source is an OGR datasource
    if geom is None:
        geom = load_datasource(source)

    if geom:
        try:
            # get the first and maybe only geometry
            check_wgs84_srs(geom[0])
        except ValueError, e:
            log.error(e)
            return None
        if rtree:
            return LimitRTreeGeometry(geom)
        else:
            log.info('You should install RTree for large --limit-to polygons')
            return LimitPolygonGeometry(build_multipolygon(geom)[1])
    return None

def check_wgs84_srs(geom):
    bbox = geom.bounds
    if bbox[0] >= -180 and bbox[1] >= -90 and bbox[2] <= 180 and bbox[3] <= 90:
        return True
    raise ValueError('Geometry is not in EPSG:4326')

class EmtpyGeometryError(Exception):
    pass

class LimitPolygonGeometry(object):
    def __init__(self, shapely_geom):
        self._geom = shapely_geom
        self._prepared_geom = None
        self._prepared_counter = 0
        self._prepared_max = 100000

    @property
    def geom(self):
        # GEOS internal data structure for prepared geometries grows over time,
        # recreate to limit memory consumption
        if not self._prepared_geom or self._prepared_counter > self._prepared_max:
            print 'create prepared'
            self._prepared_geom = shapely.prepared.prep(self._geom)
            self._prepared_counter = 0
        self._prepared_counter += 1
        return self._prepared_geom

    def intersection(self, geom):
        if self.geom.contains_properly(geom):
            # no need to limit contained geometries
            return geom

        new_geom = None
        if self.geom.intersects(geom):
            try:
                # can not use intersection with prepared geom
                new_geom = self._geom.intersection(geom)
            except TopologicalError:
                pass

        if not new_geom or new_geom.is_empty:
            raise EmtpyGeometryError('No intersection or empty geometry')

        new_geom = filter_geometry_by_type(new_geom, geom.type)
        if new_geom:
            return new_geom

        raise EmtpyGeometryError('No intersection or empty geometry')

def filter_geometry_by_type(geometry, geom_type):
    """
    Filter (multi)geometry for compatible `geom_type`,
    because we can't insert points into linestring tables for example
    """
    if geometry.type == geom_type:
        # same type is fine
        return geometry

    if geometry.type == 'MultiPolygon' and geom_type == 'Polygon':
        # polygon mappings should also support multipolygons
        return geometry

    if hasattr(geometry, 'geoms'):
        # GeometryCollection or MultiLineString? return list of geometries
        geoms = []
        for part in geometry.geoms:
            # only parts with same type
            if part.type == geom_type:
                geoms.append(part)

        if geoms:
            return geoms

    return None

def flatten_polygons(polygons):
    for polygon in polygons:
        if polygon.type == 'MultiPolygon':
            for p in polygon.geoms:
                yield p
        else:
            yield polygon

def flatten_linestrings(linestrings):
    for linestring in linestrings:
        if linestring.type == 'MultiLineString':
            for ls in linestring.geoms:
                yield ls
        else:
            yield linestring

class LimitRTreeGeometry(object):
    def __init__(self, polygons):
        index = rtree.index.Index()
        sub_polygons = []
        part_idx = 0
        for polygon in polygons:
            for part in split_polygon_at_grid(polygon):
                sub_polygons.append(part)
                index.insert(part_idx, part.bounds)
                part_idx += 1

        self.polygons = sub_polygons
        self.index = index

    def intersection(self, geom):
        intersection_ids = list(self.index.intersection(geom.bounds))

        if not intersection_ids:
            raise EmtpyGeometryError('No intersection or empty geometry')

        intersections = []
        for i in intersection_ids:

            polygon = self.polygons[i]

            if polygon.contains(geom):
                return geom

            if polygon.intersects(geom):
                try:
                    new_geom_part = polygon.intersection(geom)
                    new_geom_part = filter_geometry_by_type(new_geom_part, geom.type)
                    if new_geom_part:
                        if len(intersection_ids) == 1:
                            return new_geom_part

                        if isinstance(new_geom_part, list):
                            intersections.extend(new_geom_part)
                        else:
                            intersections.append(new_geom_part)
                except TopologicalError:
                        pass

        if not intersections:
            raise EmtpyGeometryError('No intersection or empty geometry')

        # intersections from multiple sub-polygons
        # try to merge them back to a single geometry
        if geom.type.endswith('Polygon'):
            union = cascaded_union(list(flatten_polygons(intersections)))
        elif geom.type.endswith('LineString'):
            union = linemerge(list(flatten_linestrings(intersections)))
            if union.type == 'MultiLineString':
                union = list(union.geoms)
        elif geom.type == 'Point':
            union = intersections[0]
        else:
            log.warn('unexpexted geometry type %s', geom.type)
            raise EmtpyGeometryError()
        return union
<|endoftext|>


<commit_before><commit_msg>http调试信息修改<commit_after><|endoftext|>


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

"""
The analyser is the central point of the yoda project. It analyses the given script and store the results in a mongodb
collection. Configuration of the mongodb is done in the settings.py file.
.. todo:: Use a better way to store settings in order to let the user change things
`todo`

This script can't be run directly and must be imported your own script in order to analyse it. Be sure to
add the ``set_trace()`` function at the beginning and ``set_quit()`` and the end of the section of code you want to
analyse.

:Example:

import yoda.analyser

yoda.analyser.db.set_trace()

a = 1
b = 3

for c in range(10):
    d = c
    print(c)
yoda.analyser.db.set_quit()
"""

import bdb
from collections import defaultdict
from datetime import datetime
import subprocess
import sys
import timeit
import copy
from pprint import pprint
import inspect

import mongoengine

import yoda.settings as settings
from yoda.docdef import *

IS_PYTHON_3 = (sys.version_info[0] == 3)

class Yoda(bdb.Bdb):
    run = 0
    json_results = None
    instrumented_types = (int, float, str, list, dict)
    simple_types = (int, float, str)
    #instrumented_types = (dict, bytes, bool, float, int, list, object, str, tuple)
    prev_lineno = defaultdict(int)
    prev_lineno['<module>'] = 0 # Fix to 0 because there is no previous line number
    cur_framename = '<module>' # Fixed to this value because it's always the main frame name
    file_name = None
    file_id = None # File id of the saved object
    total_linenb = 0
    next_backup = 1000

    def __init__(self):
        bdb.Bdb.__init__(self)
        if settings.DEBUG is False:
            mongoengine.connect(settings.MONGODB)
        self._clear_cache()

    def _clear_cache(self):
        self.json_results = defaultdict(lambda: defaultdict(defaultdict))

    def _filter_locals(self, local_vars):
        new_locals = {}
        for name, value in list(local_vars.items()):
            if name.startswith('__'):
                continue
            if not isinstance(value, self.instrumented_types):
                continue

            new_locals[name] = [copy.deepcopy(value)]

        return new_locals

    def _get_git_revision_short_hash(self):
        """
        Get the git revision number
        :return: the last revision number or None if the file is not in a git repo
        """

        try:
            revision = subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD'])
        except:
            revision = None

        return revision

    def _get_git_username(self):
        """
        Get the git username
        :return: git username or None if not configured
        """

        try:
            username = subprocess.check_output(['git', 'config', 'user.name'])
        except:
            username = None

        return username

    def _create_new_file(self):
        for module_file, frames in self.json_results.items():
            if 'analyser.py' not in module_file:
                file = open(module_file, 'r')
                if IS_PYTHON_3:
                    file_content = file.read()
                else:
                    file_content = "Sorry but Python 2.7 is a real mess with encodings"
                file.close()

                if file_content:
                    item = File(user=self._get_git_username(), revision=self._get_git_revision_short_hash(), filename=module_file, timestamp=datetime.now(), content=file_content)
                    for name, lines in sorted(frames.items()):
                        frame = Frame(name=name)
                        for lineno, data in sorted(lines.items()):
                            line = Line(lineno = lineno, data = data)
                            frame.lines.append(line)

                        item.frames.append(frame)

                    item.save()
                    self.file_id = item.id

    def _update_file(self):
        for module_file, frames in self.json_results.items():
            if 'analyser.py' not in module_file:
                for name, lines in sorted(frames.items()):
                    for lineno, data in sorted(lines.items()):
                        line = Line(lineno = lineno, data = data)
                        item = File.objects(id=self.file_id, frames__name=name).update(push__frames__S__lines=line)

    def _populate_db(self):
        if self.file_id is None:
            self._create_new_file()
            self._clear_cache()
        else:
            self._update_file()
            self._clear_cache()

    def user_call(self, frame, args):
        self.interaction(frame, 'call', None)

    def user_line(self, frame):
        self.interaction(frame, 'line', None)

    def user_return(self, frame, value):
        self.interaction(frame, 'return', None)

    def user_exception(self, frame, exception):
        self.interaction(frame, 'line', exception)

    def interaction(self, frame, event, exception):
        if event == 'call':
            self.cur_framename = str(frame.f_code.co_name)
            self.prev_lineno[self.cur_framename] = frame.f_lineno
            self.set_step() # continue

        if event == 'line':
            cur_lineno = frame.f_lineno
            lineno = self.prev_lineno[self.cur_framename]

            locals = self._filter_locals(frame.f_locals)
            filename = frame.f_globals['__file__']

            print(inspect.getsourcelines(frame.f_code))

            if locals:
                if not self.json_results[filename][self.cur_framename].get(lineno):
                    self.json_results[filename][self.cur_framename][lineno] = locals
                else:
                    for k,v in locals.items():
                        if not self.json_results[filename][self.cur_framename][lineno].get(k):
                            self.json_results[filename][self.cur_framename][lineno][k] = v
                        else:
                            self.json_results[filename][self.cur_framename][lineno][k].append(v[0])

            self.prev_lineno[self.cur_framename] = cur_lineno
            self.total_linenb += 1
            if self.total_linenb > self.next_backup:
                self._populate_db()
                self.next_backup += self.next_backup

            self.set_step()

        if event == 'return':
            self.cur_framename = '<module>'
            self.set_step()  # continue

        if event == 'exception':
            name = frame.f_code.co_name or "<unknown>"
            print("exception in", name, exception)
            self.set_continue()  # continue

    def set_quit(self):
        self.stopframe = self.botframe
        self.returnframe = None
        self.quitting = True
        sys.settrace(None)

        if self.json_results:
            if settings.DEBUG:
                print(self.json_results)
            else:
                self._populate_db()
        print(timeit.timeit('"-".join(str(n) for n in range(100))', number=10000))

db = Yoda()
<commit_msg>Correction of the update_db function<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
The analyser is the central point of the yoda project. It analyses the given script and store the results in a mongodb
collection. Configuration of the mongodb is done in the settings.py file.
.. todo:: Use a better way to store settings in order to let the user change things
`todo`

This script can't be run directly and must be imported your own script in order to analyse it. Be sure to
add the ``set_trace()`` function at the beginning and ``set_quit()`` and the end of the section of code you want to
analyse.

:Example:

import yoda.analyser

yoda.analyser.db.set_trace()

a = 1
b = 3

for c in range(10):
    d = c
    print(c)
yoda.analyser.db.set_quit()
"""

import bdb
from collections import defaultdict
from datetime import datetime
import subprocess
import sys
import timeit
import copy
from pprint import pprint
import inspect

import mongoengine

import yoda.settings as settings
from yoda.docdef import *

IS_PYTHON_3 = (sys.version_info[0] == 3)

class Yoda(bdb.Bdb):
    run = 0
    json_results = None
    instrumented_types = (int, float, str, list, dict)
    simple_types = (int, float, str)
    #instrumented_types = (dict, bytes, bool, float, int, list, object, str, tuple)
    prev_lineno = defaultdict(int)
    prev_lineno['<module>'] = 0 # Fix to 0 because there is no previous line number
    cur_framename = '<module>' # Fixed to this value because it's always the main frame name
    file_name = None
    file_id = None # File id of the saved object
    total_linenb = 0
    next_backup = 10

    def __init__(self):
        bdb.Bdb.__init__(self)
        if settings.DEBUG is False:
            mongoengine.connect(settings.MONGODB)
        self._clear_cache()

    def _clear_cache(self):
        self.json_results = defaultdict(lambda: defaultdict(defaultdict))

    def _filter_locals(self, local_vars):
        new_locals = {}
        for name, value in list(local_vars.items()):
            if name.startswith('__'):
                continue
            if not isinstance(value, self.instrumented_types):
                continue

            new_locals[name] = [copy.deepcopy(value)]

        return new_locals

    def _get_git_revision_short_hash(self):
        """
        Get the git revision number
        :return: the last revision number or None if the file is not in a git repo
        """

        try:
            revision = subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD'])
        except:
            revision = None

        return revision

    def _get_git_username(self):
        """
        Get the git username
        :return: git username or None if not configured
        """

        try:
            username = subprocess.check_output(['git', 'config', 'user.name'])
        except:
            username = None

        return username

    def _create_new_file(self):
        for module_file, frames in self.json_results.items():
            if 'analyser.py' not in module_file:
                file = open(module_file, 'r')
                if IS_PYTHON_3:
                    file_content = file.read()
                else:
                    file_content = "Sorry but Python 2.7 is a real mess with encodings"
                file.close()

                if file_content:
                    item = File(user=self._get_git_username(), revision=self._get_git_revision_short_hash(), filename=module_file, timestamp=datetime.now(), content=file_content)
                    for name, lines in sorted(frames.items()):
                        frame = Frame(name=name)
                        for lineno, data in sorted(lines.items()):
                            line = Line(lineno = lineno, data = data)
                            frame.lines.append(line)

                        item.frames.append(frame)

                    item.save()
                    self.file_id = item.id

    def _update_file(self):
        for module_file, frames in self.json_results.items():
            if 'analyser.py' not in module_file:
                for name, lines in sorted(frames.items()):
                    if not File.objects(id=self.file_id, frames__name=name):
                        frame = Frame(name=name)
                        File.objects(id=self.file_id).update(push__frames=frame)
                    for lineno, data in sorted(lines.items()):
                        line = Line(lineno = lineno, data = data)
                        File.objects(id=self.file_id, frames__name=name).update(push__frames__S__lines=line)

    def _populate_db(self):
        if self.file_id is None:
            self._create_new_file()
            self._clear_cache()
            print("FILE CREATED")
        else:
            self._update_file()
            self._clear_cache()
            print("FILE UPDATED")

    def user_call(self, frame, args):
        self.interaction(frame, 'call', None)

    def user_line(self, frame):
        self.interaction(frame, 'line', None)

    def user_return(self, frame, value):
        self.interaction(frame, 'return', None)

    def user_exception(self, frame, exception):
        self.interaction(frame, 'line', exception)

    def interaction(self, frame, event, exception):
        if self.file_name is None:
            self.file_name = inspect.getfile(frame)
            print(self.file_name)

        if self.file_name is inspect.getfile(frame):

            if event == 'call':
                self.cur_framename = str(frame.f_code.co_name)
                self.prev_lineno[self.cur_framename] = frame.f_lineno
                self.set_step() # continue

            if event == 'line':
                cur_lineno = frame.f_lineno
                lineno = self.prev_lineno[self.cur_framename]

                locals = self._filter_locals(frame.f_locals)
                filename = frame.f_globals['__file__']

                if locals:
                    if not self.json_results[filename][self.cur_framename].get(lineno):
                        self.json_results[filename][self.cur_framename][lineno] = locals
                    else:
                        for k,v in locals.items():
                            if not self.json_results[filename][self.cur_framename][lineno].get(k):
                                self.json_results[filename][self.cur_framename][lineno][k] = v
                            else:
                                self.json_results[filename][self.cur_framename][lineno][k].append(v[0])

                self.prev_lineno[self.cur_framename] = cur_lineno
                self.total_linenb += 1
                if self.total_linenb > self.next_backup:
                    self._populate_db()
                    self.next_backup += self.next_backup

                self.set_step()

            if event == 'return':
                self.cur_framename = '<module>'
                self.set_step()  # continue

            if event == 'exception':
                name = frame.f_code.co_name or "<unknown>"
                print("exception in", name, exception)
                self.set_continue()  # continue

    def set_quit(self):
        self.stopframe = self.botframe
        self.returnframe = None
        self.quitting = True
        sys.settrace(None)

        if self.json_results:
            if settings.DEBUG:
                print(self.json_results)
            else:
                self._populate_db()
        print(timeit.timeit('"-".join(str(n) for n in range(100))', number=10000))

db = Yoda()
<|endoftext|>


<commit_before>"""Next-gen alignments with TopHat a spliced read mapper for RNA-seq experiments.

http://tophat.cbcb.umd.edu
"""
import sh
import os
import shutil
import subprocess
from contextlib import closing
import glob

import pysam
import numpy

from bcbio.pipeline import config_utils
from bcbio.ngsalign import bowtie, bowtie2
from bcbio.utils import safe_makedir, file_exists
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger

galaxy_location_file = "bowtie_indices.loc"

_out_fnames = ["accepted_hits.sam", "junctions.bed",
               "insertions.bed", "deletions.bed"]

def tophat_align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
                 rg_name=None):
    """
    run alignment using Tophat v1
    """
    if _ref_version(ref_file) != 1:
        logger.error("Tophat v1 needs Bowtie version 1 libraries "
                     "(the extension should be .ebwt). Download "
                     "v1 libraries or build new ones with bowtie-build.")
        exit(1)

    qual_format = config["algorithm"].get("quality_format", None)

    if qual_format is None or qual_format.lower() == "illumina":
        qual_flags = ["--solexa1.3-quals"]
    elif qual_format == "solexa":
        qual_flags = ["--solexa-quals"]
    else:
        qual_flags = []

    gtf_file = config.get("gtf", None)
    if gtf_file is not None:
        gtf_flags = ["--GTF", gtf_file]
    else:
        gtf_flags = []

    cores = config.get("resources", {}).get("tophat", {}).get("cores", None)
    core_flags = ["-p", str(cores)] if cores else []
    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])
    files = [ref_file, fastq_file]
    if not file_exists(out_file):
        with file_transaction(out_dir) as tx_out_dir:
            safe_makedir(tx_out_dir)
            cl = [config_utils.get_program("tophat", config)]
            cl += core_flags
            cl += qual_flags
            cl += gtf_flags
            cl += ["-m", str(config["algorithm"].get("max_errors", 0)),
                   "--output-dir", tx_out_dir,
                   "--no-convert-bam"]
            if pair_file:
                d, d_stdev = _estimate_paired_innerdist(fastq_file, pair_file, ref_file,
                                                        out_base, tx_out_dir, config)
                cl += ["--mate-inner-dist", str(d),
                       "--mate-std-dev", str(d_stdev)]
                files.append(pair_file)
            cl += files
            child = subprocess.check_call(cl)
    out_file_final = os.path.join(out_dir, "%s.sam" % out_base)
    if not os.path.exists(out_file_final):
        os.symlink(out_file, out_file_final)
    return out_file_final

def tophat2_align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
                 rg_name=None):
    """
    run alignment using Tophat v2
    """

    """
    select the correct bowtie to run based on the reference version
    """

    qual_format = config["algorithm"].get("quality_format", None)

    if qual_format is None or qual_format.lower() == "illumina":
        qual_flags = ["--solexa1.3-quals"]
    elif qual_format == "solexa":
        qual_flags = ["--solexa-quals"]
    else:
        qual_flags = []

    gtf_file = config.get("gtf", None)
    if gtf_file is not None:
        gtf_flags = ["--GTF", gtf_file]
    else:
        gtf_flags = []

    if _ref_version(ref_file) == 1:
        ref_flags = ["--bowtie1"]
    else:
        ref_flags = []

    cores = config.get("resources", {}).get("tophat", {}).get("cores", None)
    core_flags = ["-p", str(cores)] if cores else []
    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])
    files = [ref_file, fastq_file]
    if not file_exists(out_file):
        with file_transaction(out_dir) as tx_out_dir:
            safe_makedir(tx_out_dir)
            cl = [config_utils.get_program("tophat", config)]
            cl += core_flags
            cl += qual_flags
            cl += gtf_flags
            cl += ref_flags
            cl += ["-m", str(config["algorithm"].get("max_errors", 0)),
                   "--output-dir", tx_out_dir,
                   "--no-convert-bam"]
            if pair_file:
                d, d_stdev = _estimate_paired_innerdist(fastq_file, pair_file, ref_file,
                                                        out_base, tx_out_dir, config)
                cl += ["--mate-inner-dist", str(d),
                       "--mate-std-dev", str(d_stdev)]
                files.append(pair_file)
            cl += files
            child = subprocess.check_call(cl)
    out_file_final = os.path.join(out_dir, "%s.sam" % out_base)
    if not os.path.exists(out_file_final):
        os.symlink(out_file, out_file_final)
    return out_file_final

def align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
          rg_name=None):

    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])

    if file_exists(out_file):
        return out_file

    if not _bowtie_ref_match(ref_file, config):
        logger.error("Bowtie version %d was detected but the reference "
                     "file %s is built for version %d. Download version "
                     "%d or build it with bowtie-build."
                     % (_bowtie_major_version(config), ref_file,
                        _ref_version(ref_file),
                        _bowtie_major_version(config)))
        exit(1)

    tophat_version = _tophat_major_version(config)
    if tophat_version == 1:
        out_files = tophat_align(fastq_file, pair_file, ref_file, out_base,
                                 align_dir, config, rg_name=None)
    elif tophat_version == 2:
        out_files = tophat2_align(fastq_file, pair_file, ref_file, out_base,
                                  align_dir, config, rg_name=None)
    else:
        raise ValueError("Tophat major version %d is not supported yet."
                         % (tophat_version))

    return out_files

def _estimate_paired_innerdist(fastq_file, pair_file, ref_file, out_base,
                               out_dir, config):
    """Use Bowtie to estimate the inner distance of paired reads.
    """
    # skip initial reads for large file, but not for smaller
    dists = _bowtie_for_innerdist("1000000", fastq_file, pair_file, ref_file,
                                  out_base, out_dir, config)
    if len(dists) == 0:
        dists = _bowtie_for_innerdist("1", fastq_file, pair_file, ref_file,
                                      out_base, out_dir, config, True)
    return int(round(numpy.mean(dists))), int(round(numpy.std(dists)))

def _bowtie_for_innerdist(start, fastq_file, pair_file, ref_file, out_base,
                          out_dir, config, remove_workdir=False):
    work_dir = os.path.join(out_dir, "innerdist_estimate")
    if os.path.exists(work_dir):
        shutil.rmtree(work_dir)
    safe_makedir(work_dir)
    extra_args = ["-s", str(start), "-u", "250000"]
    bowtie_runner = _select_bowtie_version(config)
    out_sam = bowtie_runner.align(fastq_file, pair_file, ref_file, out_base,
                                  work_dir, config, extra_args)
    dists = []
    with closing(pysam.Samfile(out_sam)) as work_sam:
        for read in work_sam:
            if not read.is_unmapped and read.is_read1:
                dists.append(abs(read.isize) - 2 * read.rlen)
    return dists

def _bowtie_major_version(config):
    bowtie_runner = sh.Command(config_utils.get_program("bowtie", config,
                                                        default="tophat"))
    """
    bowtie --version returns strings like this:
    bowtie version 0.12.7
    32-bit
    Built on Franklin.local
    Tue Sep  7 14:25:02 PDT 2010
    """
    version_line = str(bowtie_runner(version=True)).split("\n")[0]
    version_string = version_line.strip().split()[2]
    major_version = int(version_string.split(".")[0])
    # bowtie version 1 has a leading character of 0
    if major_version == 0:
        major_version += 1
    return major_version

def _tophat_major_version(config):
    tophat_runner = sh.Command(config_utils.get_program("tophat", config,
                                                        default="tophat"))

    # tophat --version returns strings like this: Tophat v2.0.4
    version_string = str(tophat_runner(version=True)).strip().split()[1]
    major_version = int(version_string.split(".")[0][1:])
    return major_version

def _bowtie_ref_match(ref_file, config):
    return _ref_version(ref_file) == _bowtie_major_version(config)

def _select_bowtie_version(config):
    if _bowtie_major_version(config) == 1:
        return bowtie
    else:
        return bowtie2

def _ref_version(ref_file):
    _, ext = os.path.splitext(glob.glob(ref_file + "*")[0])
    if ext == ".ebwt":
        return 1
    elif ext == ".bt2":
        return 2
    else:
        logger.error("Cannot detect which reference version %s is. "
                     "Should end in either .ebwt (bowtie) or .bt2 "
                     "(bowtie2)." % (ref_file))
        exit(1)
<commit_msg>Added reading option list from tophat resources.<commit_after>"""Next-gen alignments with TopHat a spliced read mapper for RNA-seq experiments.

http://tophat.cbcb.umd.edu
"""
import sh
import os
import shutil
import subprocess
from contextlib import closing
import glob
import pysam
import numpy
from bcbio.pipeline import config_utils
from bcbio.ngsalign import bowtie, bowtie2
from bcbio.utils import safe_makedir, file_exists, get_in, flatten
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger

galaxy_location_file = "bowtie_indices.loc"

_out_fnames = ["accepted_hits.sam", "junctions.bed",
               "insertions.bed", "deletions.bed"]

def tophat_align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
                 rg_name=None):
    """
    run alignment using Tophat v1
    """
    if _ref_version(ref_file) != 1:
        logger.error("Tophat v1 needs Bowtie version 1 libraries "
                     "(the extension should be .ebwt). Download "
                     "v1 libraries or build new ones with bowtie-build.")
        exit(1)

    qual_format = config["algorithm"].get("quality_format", None)

    if qual_format is None or qual_format.lower() == "illumina":
        qual_flags = ["--solexa1.3-quals"]
    elif qual_format == "solexa":
        qual_flags = ["--solexa-quals"]
    else:
        qual_flags = []

    gtf_file = config.get("gtf", None)
    if gtf_file is not None:
        gtf_flags = ["--GTF", gtf_file]
    else:
        gtf_flags = []

    cores = get_in(config, ("resources", "tophat", "cores"), {})
    options = list(flatten(get_in(config, ("resources", "tophat", "options"), {}).items()))
    core_flags = ["-p", str(cores)] if cores else []
    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])
    files = [ref_file, fastq_file]
    if not file_exists(out_file):
        with file_transaction(out_dir) as tx_out_dir:
            safe_makedir(tx_out_dir)
            cl = [config_utils.get_program("tophat", config)]
            cl += core_flags
            cl += qual_flags
            cl += gtf_flags
            cl += options
            cl += ["-m", str(config["algorithm"].get("max_errors", 0)),
                   "--output-dir", tx_out_dir,
                   "--no-convert-bam"]
            if pair_file:
                d, d_stdev = _estimate_paired_innerdist(fastq_file, pair_file, ref_file,
                                                        out_base, tx_out_dir, config)
                cl += ["--mate-inner-dist", str(d),
                       "--mate-std-dev", str(d_stdev)]
                files.append(pair_file)
            cl += files
            child = subprocess.check_call(cl)
    out_file_final = os.path.join(out_dir, "%s.sam" % out_base)
    if not os.path.exists(out_file_final):
        os.symlink(out_file, out_file_final)
    return out_file_final

def tophat2_align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
                 rg_name=None):
    """
    run alignment using Tophat v2
    """

    """
    select the correct bowtie to run based on the reference version
    """

    qual_format = config["algorithm"].get("quality_format", None)

    if qual_format is None or qual_format.lower() == "illumina":
        qual_flags = ["--solexa1.3-quals"]
    elif qual_format == "solexa":
        qual_flags = ["--solexa-quals"]
    else:
        qual_flags = []

    gtf_file = config.get("gtf", None)
    if gtf_file is not None:
        gtf_flags = ["--GTF", gtf_file]
    else:
        gtf_flags = []

    if _ref_version(ref_file) == 1:
        ref_flags = ["--bowtie1"]
    else:
        ref_flags = []

    cores = config.get("resources", {}).get("tophat", {}).get("cores", None)
    core_flags = ["-p", str(cores)] if cores else []
    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])
    options = list(flatten(get_in(config, ("resources", "tophat", "options"), {}).items()))
    files = [ref_file, fastq_file]
    if not file_exists(out_file):
        with file_transaction(out_dir) as tx_out_dir:
            safe_makedir(tx_out_dir)
            cl = [config_utils.get_program("tophat", config)]
            cl += core_flags
            cl += qual_flags
            cl += gtf_flags
            cl += ref_flags
            cl += options
            cl += ["-m", str(config["algorithm"].get("max_errors", 0)),
                   "--output-dir", tx_out_dir,
                   "--no-convert-bam"]
            if pair_file:
                d, d_stdev = _estimate_paired_innerdist(fastq_file, pair_file, ref_file,
                                                        out_base, tx_out_dir, config)
                cl += ["--mate-inner-dist", str(d),
                       "--mate-std-dev", str(d_stdev)]
                files.append(pair_file)
            cl += files
            child = subprocess.check_call(cl)
    out_file_final = os.path.join(out_dir, "%s.sam" % out_base)
    if not os.path.exists(out_file_final):
        os.symlink(out_file, out_file_final)
    return out_file_final

def align(fastq_file, pair_file, ref_file, out_base, align_dir, config,
          rg_name=None):

    out_dir = os.path.join(align_dir, "%s_tophat" % out_base)
    out_file = os.path.join(out_dir, _out_fnames[0])

    if file_exists(out_file):
        return out_file

    if not _bowtie_ref_match(ref_file, config):
        logger.error("Bowtie version %d was detected but the reference "
                     "file %s is built for version %d. Download version "
                     "%d or build it with bowtie-build."
                     % (_bowtie_major_version(config), ref_file,
                        _ref_version(ref_file),
                        _bowtie_major_version(config)))
        exit(1)

    tophat_version = _tophat_major_version(config)
    if tophat_version == 1:
        out_files = tophat_align(fastq_file, pair_file, ref_file, out_base,
                                 align_dir, config, rg_name=None)
    elif tophat_version == 2:
        out_files = tophat2_align(fastq_file, pair_file, ref_file, out_base,
                                  align_dir, config, rg_name=None)
    else:
        raise ValueError("Tophat major version %d is not supported yet."
                         % (tophat_version))

    return out_files

def _estimate_paired_innerdist(fastq_file, pair_file, ref_file, out_base,
                               out_dir, config):
    """Use Bowtie to estimate the inner distance of paired reads.
    """
    # skip initial reads for large file, but not for smaller
    dists = _bowtie_for_innerdist("1000000", fastq_file, pair_file, ref_file,
                                  out_base, out_dir, config)
    if len(dists) == 0:
        dists = _bowtie_for_innerdist("1", fastq_file, pair_file, ref_file,
                                      out_base, out_dir, config, True)
    return int(round(numpy.mean(dists))), int(round(numpy.std(dists)))

def _bowtie_for_innerdist(start, fastq_file, pair_file, ref_file, out_base,
                          out_dir, config, remove_workdir=False):
    work_dir = os.path.join(out_dir, "innerdist_estimate")
    if os.path.exists(work_dir):
        shutil.rmtree(work_dir)
    safe_makedir(work_dir)
    extra_args = ["-s", str(start), "-u", "250000"]
    bowtie_runner = _select_bowtie_version(config)
    out_sam = bowtie_runner.align(fastq_file, pair_file, ref_file, out_base,
                                  work_dir, config, extra_args)
    dists = []
    with closing(pysam.Samfile(out_sam)) as work_sam:
        for read in work_sam:
            if not read.is_unmapped and read.is_read1:
                dists.append(abs(read.isize) - 2 * read.rlen)
    return dists

def _bowtie_major_version(config):
    bowtie_runner = sh.Command(config_utils.get_program("bowtie", config,
                                                        default="tophat"))
    """
    bowtie --version returns strings like this:
    bowtie version 0.12.7
    32-bit
    Built on Franklin.local
    Tue Sep  7 14:25:02 PDT 2010
    """
    version_line = str(bowtie_runner(version=True)).split("\n")[0]
    version_string = version_line.strip().split()[2]
    major_version = int(version_string.split(".")[0])
    # bowtie version 1 has a leading character of 0
    if major_version == 0:
        major_version += 1
    return major_version

def _tophat_major_version(config):
    tophat_runner = sh.Command(config_utils.get_program("tophat", config,
                                                        default="tophat"))

    # tophat --version returns strings like this: Tophat v2.0.4
    version_string = str(tophat_runner(version=True)).strip().split()[1]
    major_version = int(version_string.split(".")[0][1:])
    return major_version

def _bowtie_ref_match(ref_file, config):
    return _ref_version(ref_file) == _bowtie_major_version(config)

def _select_bowtie_version(config):
    if _bowtie_major_version(config) == 1:
        return bowtie
    else:
        return bowtie2

def _ref_version(ref_file):
    _, ext = os.path.splitext(glob.glob(ref_file + "*")[0])
    if ext == ".ebwt":
        return 1
    elif ext == ".bt2":
        return 2
    else:
        logger.error("Cannot detect which reference version %s is. "
                     "Should end in either .ebwt (bowtie) or .bt2 "
                     "(bowtie2)." % (ref_file))
        exit(1)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
States to manage OpenStack Nova
===============================

.. versionadded:: Helium

Use this minion to do things with nova:

.. code-block:: yaml

    myblock:
      nova.volume_exists:
        - size: 100
'''

# Import salt cloud utils for check
import salt.utils.cloud as suc

def __virtual__():
    '''
    Only load if the nova module is available in __salt__
    '''
    return 'nova' if 'nova.server_list' in __salt__ else False

def _check_name(name):
    ret = {'name': name,
           'changes': {},
           'result': None,
           'comment': ''}
    if suc.check_name(name, 'a-zA-Z0-9._-'):
        ret['comment'] = 'Invalid characters in name.'
        ret['result'] = False
        return ret
    else:
        ret['result'] = True
        return ret

def volume_exists(name, profile=None, **kwargs):
    '''
    Check that a block volume exists.
    '''
    ret = _check_name(name)
    if not ret['result']:
        return ret

    volume = __salt__['nova.volume_show'](name=name, profile=profile)

    if volume:
        ret['comment'] = 'Volume exists: {0}'.format(name)
        ret['result'] = True
        return ret
    elif __opts__['test']:
        ret['comment'] = 'Volume {0} will be created.'.format(name)
        ret['result'] = None
        return ret

    response = __salt__['nova.volume_create'](
        name=name,
        profile=profile,
        **kwargs
    )
    if response:
        ret['result'] = True
        ret['comment'] = 'Volume {0} was created'.format(name)
        ret['changes'] = {'old': None, 'new': response}
    else:
        ret['result'] = False
        ret['comment'] = 'Volume {0} failed to create.'.format(name)
    return ret

def volume_attached(name, server_name, profile=None, **kwargs):
    '''
    Check if a block volume is attached.
    '''
    ret = _check_name(name)
    if not ret['result']:
        return ret

    ret = _check_name(server_name)
    if not ret['result']:
        return ret

    volume = __salt__['nova.volume_show'](
        name=name,
        profile=profile
    )
    server = __salt__['nova.server_by_name'](server_name, profile=profile)

    if volume and volume['attachments']:
        ret['comment'] = ('Volume {name} is already'
                          'attached: {attachments}').format(**volume)
        ret['result'] = True
        return ret
    elif __opts__['test']:
        ret['comment'] = 'Volume {0} will be will be attached.'.format(
            name
        )
        ret['result'] = None
        return ret
    elif not volume:
        ret['comment'] = 'Volume {0} does not exist'.format(name)
        ret['result'] = False
        return ret
    elif not server:
        ret['comment'] = 'Server {0} does not exist'.format(server_name)
        ret['result'] = False
        return ret

    response = __salt__['nova.volume_attach'](
        name,
        server_name,
        profile=profile,
        **kwargs
    )
    if response:
        ret['result'] = True
        ret['comment'] = 'Volume {0} was created'.format(name)
        ret['changes'] = {'old': volume, 'new': response}
    else:
        ret['result'] = False
        ret['comment'] = 'Volume {0} failed to attach.'.format(name)
    return ret
<commit_msg>Remove nova state<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-

""" License

    Copyright (C) 2013 YunoHost

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program; if not, see http://www.gnu.org/licenses

"""

""" yunohost_hook.py

    Manage hooks
"""
import os
import sys
import re
import json
from yunohost import YunoHostError, YunoHostLDAP, win_msg, colorize

hook_folder = '/user/share/yunohost/hooks/'

def hook_add(action, file):
    """
    Store hook script to filsystem

    Keyword argument:
        file -- Script to check
        action -- Action folder to store into

    """
    try: os.listdir(hook_folder + action)
    except OSError: os.makedirs(hook_folder + action)

    os.system('cp '+ file +' '+ hook_folder + action)

def hook_callback(action):
    """
    Execute all scripts binded to an action

    Keyword argument:
        action -- Action name

    """
    with YunoHostLDAP() as yldap:
        try: os.listdir(hook_folder + action)
        except OSError: os.makedirs(hook_folder + action)

        for hook in os.listdir(hook_folder + action):
            hook_exec(file=hook_folder + action +'/'+ hook)

def hook_check(file):
    """
    Parse the script file and get arguments

    Keyword argument:
        file -- File to check

    """
    try:
        with open(file[:file.index('scripts/')] + 'manifest.json') as f:
            manifest = json.loads(str(f.read()))
    except:
        raise YunoHostError(22, _("Invalid app package"))

    action = file[file.index('scripts/') + 8:]
    if action in manifest["arguments"]:
        return manifest["arguments"][action]
    else:
        return {}

def hook_exec(file, args=None):
    """
    Execute hook from a file with arguments

    Keyword argument:
        file -- Script to execute
        args -- Arguments to pass to the script

    """
    with YunoHostLDAP() as yldap:
        required_args = hook_check(file)
        if args is None:
            args = {}

        arg_list = []
        for arg in required_args:
            if arg['name'] in args:
                if 'choices' in arg and args[arg['name']] not in arg['choices']:
                    raise YunoHostError(22, _("Invalid choice") + ': ' + args[arg['name']])
                arg_list.append(args[arg['name']])
            else:
                if os.isatty(1) and 'ask' in arg:
                    ask_string = arg['ask']['en'] #TODO: I18n
                    if 'choices' in arg:
                        ask_string = ask_string +' ('+ '|'.join(arg['choices']) +')'
                    if 'default' in arg:
                        ask_string = ask_string +' (default: '+ arg['default'] +')'
                    arg_list.append(raw_input(colorize(ask_string + ': ', 'cyan')))
                elif 'default' in arg:
                    arg_list.append(arg['default'])
                else:
                    raise YunoHostError(22, _("Missing arguments") + ': ' + arg['name'])

        file_path = "./"
        if "/" in file and file[0:2] != file_path:
            file_path = os.path.dirname(file)
            file = file.replace(file_path +"/", "")
        return os.system('su - admin -c "cd \\"'+ file_path +'\\" && bash \\"'+ file +'\\" '+ ' '.join(arg_list) +'"') #TODO: Allow python script
<commit_msg>Better argument input handling<commit_after># -*- coding: utf-8 -*-

""" License

    Copyright (C) 2013 YunoHost

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published
    by the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program; if not, see http://www.gnu.org/licenses

"""

""" yunohost_hook.py

    Manage hooks
"""
import os
import sys
import re
import json
from yunohost import YunoHostError, YunoHostLDAP, win_msg, colorize

hook_folder = '/user/share/yunohost/hooks/'

def hook_add(action, file):
    """
    Store hook script to filsystem

    Keyword argument:
        file -- Script to check
        action -- Action folder to store into

    """
    try: os.listdir(hook_folder + action)
    except OSError: os.makedirs(hook_folder + action)

    os.system('cp '+ file +' '+ hook_folder + action)

def hook_callback(action):
    """
    Execute all scripts binded to an action

    Keyword argument:
        action -- Action name

    """
    with YunoHostLDAP() as yldap:
        try: os.listdir(hook_folder + action)
        except OSError: os.makedirs(hook_folder + action)

        for hook in os.listdir(hook_folder + action):
            hook_exec(file=hook_folder + action +'/'+ hook)

def hook_check(file):
    """
    Parse the script file and get arguments

    Keyword argument:
        file -- File to check

    """
    try:
        with open(file[:file.index('scripts/')] + 'manifest.json') as f:
            manifest = json.loads(str(f.read()))
    except:
        raise YunoHostError(22, _("Invalid app package"))

    action = file[file.index('scripts/') + 8:]
    if action in manifest["arguments"]:
        return manifest["arguments"][action]
    else:
        return {}

def hook_exec(file, args=None):
    """
    Execute hook from a file with arguments

    Keyword argument:
        file -- Script to execute
        args -- Arguments to pass to the script

    """
    with YunoHostLDAP() as yldap:
        required_args = hook_check(file)
        if args is None:
            args = {}

        arg_list = []
        for arg in required_args:
            if arg['name'] in args:
                if 'choices' in arg and args[arg['name']] not in arg['choices']:
                    raise YunoHostError(22, _("Invalid choice") + ': ' + args[arg['name']])
                arg_list.append(args[arg['name']])
            else:
                if os.isatty(1) and 'ask' in arg:
                    ask_string = arg['ask']['en'] #TODO: I18n
                    if 'choices' in arg:
                        ask_string = ask_string +' ('+ '|'.join(arg['choices']) +')'
                    if 'default' in arg:
                        ask_string = ask_string +' (default: '+ arg['default'] +')'

                    input_string = raw_input(colorize(ask_string + ': ', 'cyan'))

                    if input_string == '' and 'default' in arg:
                        input_string = arg['default']

                    arg_list.append(input_string)
                elif 'default' in arg:
                    arg_list.append(arg['default'])
                else:
                    raise YunoHostError(22, _("Missing arguments") + ': ' + arg['name'])

        file_path = "./"
        if "/" in file and file[0:2] != file_path:
            file_path = os.path.dirname(file)
            file = file.replace(file_path +"/", "")
        return os.system('su - admin -c "cd \\"'+ file_path +'\\" && bash \\"'+ file +'\\" '+ ' '.join(arg_list) +'"') #TODO: Allow python script
<|endoftext|>


<commit_before>from datetime import date, timedelta
from typing import Tuple, Iterator

class Week:
    """A Week object represents a week in epidemiological week calendar."""

    __slots__ = "_year", "_week", "_method"

    def __init__(
        self, year: int, week: int, method: str = "CDC", validate: bool = True
    ) -> None:
        """
        :param year: Epidemiological year
        :type year: int
        :param week: Epidemiological week
        :type week: int
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        :param validate: Whether to validate year, week and method or not
            (default is ``True``)
        :type validate: bool
        """

        if validate:
            _check_year(year)
            _check_method(method)
            _check_week(year, week, method)

        self._year = year
        self._week = week
        self._method = method

    def __repr__(self) -> str:
        return "Week({}, {}, {})".format(self._year, self._week, self._method)

    def __str__(self) -> str:
        return self.isoformat()

    def __hash__(self) -> int:
        return hash((self.year, self.week, self.method))

    def __eq__(self, other: "Week") -> bool:
        return self._compare(other) == 0

    def __gt__(self, other: "Week") -> bool:
        return self._compare(other) > 0

    def __ge__(self, other: "Week") -> bool:
        return self._compare(other) >= 0

    def __lt__(self, other: "Week") -> bool:
        return self._compare(other) < 0

    def __le__(self, other: "Week") -> bool:
        return self._compare(other) <= 0

    def _compare(self, other: "Week") -> int:
        """Compare two Week objects after checking if they are comparable."""
        class_name = self.__class__.__name__
        other_name = type(other).__name__
        if not isinstance(other, self.__class__):
            raise TypeError(f"can't compare '{class_name}' to '{other_name}'")
        if self.method != other.method:
            raise TypeError(
                f"can't compare '{class_name}' objects with different "
                f"calculation methods"
            )
        x = self.weektuple()
        y = other.weektuple()
        return 0 if x == y else 1 if x > y else -1

    def __add__(self, other: int) -> "Week":
        if not isinstance(other, int):
            raise TypeError("second operand must be 'int'")
        new_date = self.startdate() + timedelta(weeks=other)
        return Week.fromdate(new_date, self._method)

    def __sub__(self, other: int) -> "Week":
        if not isinstance(other, int):
            raise TypeError("second operand must be 'int'")
        return self + (-other)

    def __contains__(self, other: date) -> bool:
        if not isinstance(other, date):
            raise TypeError("tested operand must be 'date' object")
        return other in self.iterdates()

    @classmethod
    def fromdate(cls, date: date, method: str = "CDC") -> "Week":
        """Construct Week object from a date.

        :param date: Date object
        :type date: datetime.date
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        _check_method(method)
        year = date.year
        date_ordinal = date.toordinal()
        year_start_ordinal = _year_start(year, method)
        week = (date_ordinal - year_start_ordinal) // 7
        if week < 0:
            year -= 1
            year_start_ordinal = _year_start(year, method)
            week = (date_ordinal - year_start_ordinal) // 7
        elif week >= 52:
            year_start_ordinal = _year_start(year + 1, method)
            if date_ordinal >= year_start_ordinal:
                year += 1
                week = 0
        week += 1
        return cls(year, week, method, validate=False)

    @classmethod
    def fromstring(cls, week_string: str, method: str = "CDC") -> "Week":
        """Construct Week object from a formatted string.

        :param week_string: Week string formatted as ‘YYYYWW’, ‘YYYYWWW’,
            or ‘YYYY-WWW’ for example ‘201908’, ‘2019W08’, or ‘2019-W08’.
            If the string ends with weekday as in ISO formats, weekday will
            be ignored.
        :type week_string: str
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        week_string = week_string.replace("-", "").replace("W", "")
        year = int(week_string[:4])
        week = int(week_string[4:6])
        return cls(year, week, method)

    @classmethod
    def thisweek(cls, method: str = "CDC") -> "Week":
        """Construct Week object from current date.

        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        return cls.fromdate(date.today(), method)

    @property
    def year(self) -> int:
        """Return year as an integer"""
        return self._year

    @property
    def week(self) -> int:
        """Return week number as an integer"""
        return self._week

    @property
    def method(self) -> str:
        """Return calculation method as a string"""
        return self._method.upper()

    def weektuple(self) -> Tuple[int, int]:
        """Return week as a tuple of (year, week)."""
        return self._year, self._week

    def cdcformat(self) -> str:
        """Return a string representing the week in CDC format ‘YYYYWW’ for
        example ‘201908’.
        """

        return f"{self.year:04}{self.week:02}"

    def isoformat(self) -> str:
        """Return a string representing the week in compact form of ISO format
        ‘YYYYWww’.
        """

        return "{:04}W{:02}".format(self._year, self._week)

    def startdate(self) -> date:
        """Return date for first day of week."""
        year_start_ordinal = _year_start(self._year, self._method)
        week_start_ordinal = year_start_ordinal + ((self._week - 1) * 7)
        startdate = date.fromordinal(week_start_ordinal)
        return startdate

    def enddate(self) -> date:
        """Return date for last day of week."""
        enddate = self.startdate() + timedelta(days=6)
        return enddate

    def iterdates(self) -> Iterator[date]:
        """Return an iterator that yield datetime.date objects for all days of
        week."""

        startdate = self.startdate()
        for day in range(0, 7):
            yield startdate + timedelta(days=day)

    def daydate(self, weekday: int = 6) -> date:
        """Return date for specific weekday of week.

        :param weekday: Week day, which may be ``0..6`` where Monday is 0 and
            Sunday is 6 (default is ``6``)
        :type weekday: int
        """

        days = (_method_adjustment(self.method) + weekday) % 7
        return self.startdate() + timedelta(days=days)

class Year:
    """A Year object represents a year in epidemiological week calendar."""

    __slots__ = "_year", "_method"

    def __init__(self, year: int, method: str = "CDC") -> None:
        """
        :param year: Epidemiological year
        :type year: int
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        _check_year(year)
        _check_method(method)
        self._year = year
        self._method = method

    def __repr__(self) -> str:
        return "Year({}, {})".format(self._year, self._method)

    def __str__(self) -> str:
        return "{:04}".format(self._year)

    def __hash__(self) -> int:
        return hash((self.year, self.method))

    @classmethod
    def thisyear(cls, method: str = "CDC") -> "Year":
        """Construct Year object from current date.

        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        return cls(date.today().year, method)

    @property
    def year(self) -> int:
        """Return year as an integer"""
        return self._year

    @property
    def method(self) -> str:
        """Return calculation method as a string"""
        return self._method.upper()

    def totalweeks(self) -> int:
        """Return number of weeks in year."""
        return _year_total_weeks(self._year, self._method)

    def startdate(self) -> date:
        """Return date for first day of first week of year."""
        year_start_ordinal = _year_start(self._year, self._method)
        return date.fromordinal(year_start_ordinal)

    def enddate(self) -> date:
        """Return date for last day of last week of year."""
        year_end_ordinal = _year_start(self._year + 1, self._method) - 1
        return date.fromordinal(year_end_ordinal)

    def iterweeks(self) -> Iterator[Week]:
        """Return an iterator that yield Week objects for all weeks of year."""
        for week in range(1, self.totalweeks() + 1):
            yield Week(self.year, week, self.method, validate=False)

def _check_year(year: int) -> None:
    """Check value of year."""
    if not 1 <= year <= 9999:
        raise ValueError("year must be in 1..9999")

def _check_week(year: int, week: int, method: str) -> None:
    """Check value of week."""
    weeks = _year_total_weeks(year, method)
    if not 1 <= week <= weeks:
        raise ValueError("week must be in 1..{} for year".format(weeks))

def _check_method(method: str) -> None:
    """Check value of calculation method."""
    methods = ["CDC", "ISO"]
    if method.upper() not in methods:
        raise ValueError(f"method must be '{methods[0]}' or '{methods[1]}'")

def _method_adjustment(method: str) -> int:
    """Return needed adjustment based on first day of week."""
    first_day = ("Mon", "Sun")
    if method.upper() == "CDC":
        return first_day.index("Sun")
    return first_day.index("Mon")

def _year_start(year: int, method: str) -> int:
    """Return ordinal for first day of first week for year."""
    adjustment = _method_adjustment(method)
    mid_weekday = 3 - adjustment  # Sun is 6 .. Mon is 0
    jan1 = date(year, 1, 1)
    jan1_ordinal = jan1.toordinal()
    jan1_weekday = jan1.weekday()
    week1_start_ordinal = jan1_ordinal - jan1_weekday - adjustment
    if jan1_weekday > mid_weekday:
        week1_start_ordinal += 7
    return week1_start_ordinal

def _year_total_weeks(year: int, method: str) -> int:
    """Return number of weeks in year."""
    year_start_ordinal = _year_start(year, method)
    next_year_start_ordinal = _year_start(year + 1, method)
    weeks = (next_year_start_ordinal - year_start_ordinal) // 7
    return weeks
<commit_msg>update __str__ method of Week object<commit_after>from datetime import date, timedelta
from typing import Tuple, Iterator

class Week:
    """A Week object represents a week in epidemiological week calendar."""

    __slots__ = "_year", "_week", "_method"

    def __init__(
        self, year: int, week: int, method: str = "CDC", validate: bool = True
    ) -> None:
        """
        :param year: Epidemiological year
        :type year: int
        :param week: Epidemiological week
        :type week: int
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        :param validate: Whether to validate year, week and method or not
            (default is ``True``)
        :type validate: bool
        """

        if validate:
            _check_year(year)
            _check_method(method)
            _check_week(year, week, method)

        self._year = year
        self._week = week
        self._method = method

    def __repr__(self) -> str:
        return "Week({}, {}, {})".format(self._year, self._week, self._method)

    def __str__(self) -> str:
        return self.cdcformat() if self.method == "CDC" else self.isoformat()

    def __hash__(self) -> int:
        return hash((self.year, self.week, self.method))

    def __eq__(self, other: "Week") -> bool:
        return self._compare(other) == 0

    def __gt__(self, other: "Week") -> bool:
        return self._compare(other) > 0

    def __ge__(self, other: "Week") -> bool:
        return self._compare(other) >= 0

    def __lt__(self, other: "Week") -> bool:
        return self._compare(other) < 0

    def __le__(self, other: "Week") -> bool:
        return self._compare(other) <= 0

    def _compare(self, other: "Week") -> int:
        """Compare two Week objects after checking if they are comparable."""
        class_name = self.__class__.__name__
        other_name = type(other).__name__
        if not isinstance(other, self.__class__):
            raise TypeError(f"can't compare '{class_name}' to '{other_name}'")
        if self.method != other.method:
            raise TypeError(
                f"can't compare '{class_name}' objects with different "
                f"calculation methods"
            )
        x = self.weektuple()
        y = other.weektuple()
        return 0 if x == y else 1 if x > y else -1

    def __add__(self, other: int) -> "Week":
        if not isinstance(other, int):
            raise TypeError("second operand must be 'int'")
        new_date = self.startdate() + timedelta(weeks=other)
        return Week.fromdate(new_date, self._method)

    def __sub__(self, other: int) -> "Week":
        if not isinstance(other, int):
            raise TypeError("second operand must be 'int'")
        return self + (-other)

    def __contains__(self, other: date) -> bool:
        if not isinstance(other, date):
            raise TypeError("tested operand must be 'date' object")
        return other in self.iterdates()

    @classmethod
    def fromdate(cls, date: date, method: str = "CDC") -> "Week":
        """Construct Week object from a date.

        :param date: Date object
        :type date: datetime.date
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        _check_method(method)
        year = date.year
        date_ordinal = date.toordinal()
        year_start_ordinal = _year_start(year, method)
        week = (date_ordinal - year_start_ordinal) // 7
        if week < 0:
            year -= 1
            year_start_ordinal = _year_start(year, method)
            week = (date_ordinal - year_start_ordinal) // 7
        elif week >= 52:
            year_start_ordinal = _year_start(year + 1, method)
            if date_ordinal >= year_start_ordinal:
                year += 1
                week = 0
        week += 1
        return cls(year, week, method, validate=False)

    @classmethod
    def fromstring(cls, week_string: str, method: str = "CDC") -> "Week":
        """Construct Week object from a formatted string.

        :param week_string: Week string formatted as ‘YYYYWW’, ‘YYYYWWW’,
            or ‘YYYY-WWW’ for example ‘201908’, ‘2019W08’, or ‘2019-W08’.
            If the string ends with weekday as in ISO formats, weekday will
            be ignored.
        :type week_string: str
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        week_string = week_string.replace("-", "").replace("W", "")
        year = int(week_string[:4])
        week = int(week_string[4:6])
        return cls(year, week, method)

    @classmethod
    def thisweek(cls, method: str = "CDC") -> "Week":
        """Construct Week object from current date.

        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        return cls.fromdate(date.today(), method)

    @property
    def year(self) -> int:
        """Return year as an integer"""
        return self._year

    @property
    def week(self) -> int:
        """Return week number as an integer"""
        return self._week

    @property
    def method(self) -> str:
        """Return calculation method as a string"""
        return self._method.upper()

    def weektuple(self) -> Tuple[int, int]:
        """Return week as a tuple of (year, week)."""
        return self._year, self._week

    def cdcformat(self) -> str:
        """Return a string representing the week in CDC format ‘YYYYWW’ for
        example ‘201908’.
        """

        return f"{self.year:04}{self.week:02}"

    def isoformat(self) -> str:
        """Return a string representing the week in compact form of ISO format
        ‘YYYYWww’.
        """

        return "{:04}W{:02}".format(self._year, self._week)

    def startdate(self) -> date:
        """Return date for first day of week."""
        year_start_ordinal = _year_start(self._year, self._method)
        week_start_ordinal = year_start_ordinal + ((self._week - 1) * 7)
        startdate = date.fromordinal(week_start_ordinal)
        return startdate

    def enddate(self) -> date:
        """Return date for last day of week."""
        enddate = self.startdate() + timedelta(days=6)
        return enddate

    def iterdates(self) -> Iterator[date]:
        """Return an iterator that yield datetime.date objects for all days of
        week."""

        startdate = self.startdate()
        for day in range(0, 7):
            yield startdate + timedelta(days=day)

    def daydate(self, weekday: int = 6) -> date:
        """Return date for specific weekday of week.

        :param weekday: Week day, which may be ``0..6`` where Monday is 0 and
            Sunday is 6 (default is ``6``)
        :type weekday: int
        """

        days = (_method_adjustment(self.method) + weekday) % 7
        return self.startdate() + timedelta(days=days)

class Year:
    """A Year object represents a year in epidemiological week calendar."""

    __slots__ = "_year", "_method"

    def __init__(self, year: int, method: str = "CDC") -> None:
        """
        :param year: Epidemiological year
        :type year: int
        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        _check_year(year)
        _check_method(method)
        self._year = year
        self._method = method

    def __repr__(self) -> str:
        return "Year({}, {})".format(self._year, self._method)

    def __str__(self) -> str:
        return "{:04}".format(self._year)

    def __hash__(self) -> int:
        return hash((self.year, self.method))

    @classmethod
    def thisyear(cls, method: str = "CDC") -> "Year":
        """Construct Year object from current date.

        :param method: Calculation method, which may be ``CDC`` where the week
            starts on Sunday or ``ISO`` where the week starts on Monday
            (default is ``CDC``)
        :type method: str
        """

        return cls(date.today().year, method)

    @property
    def year(self) -> int:
        """Return year as an integer"""
        return self._year

    @property
    def method(self) -> str:
        """Return calculation method as a string"""
        return self._method.upper()

    def totalweeks(self) -> int:
        """Return number of weeks in year."""
        return _year_total_weeks(self._year, self._method)

    def startdate(self) -> date:
        """Return date for first day of first week of year."""
        year_start_ordinal = _year_start(self._year, self._method)
        return date.fromordinal(year_start_ordinal)

    def enddate(self) -> date:
        """Return date for last day of last week of year."""
        year_end_ordinal = _year_start(self._year + 1, self._method) - 1
        return date.fromordinal(year_end_ordinal)

    def iterweeks(self) -> Iterator[Week]:
        """Return an iterator that yield Week objects for all weeks of year."""
        for week in range(1, self.totalweeks() + 1):
            yield Week(self.year, week, self.method, validate=False)

def _check_year(year: int) -> None:
    """Check value of year."""
    if not 1 <= year <= 9999:
        raise ValueError("year must be in 1..9999")

def _check_week(year: int, week: int, method: str) -> None:
    """Check value of week."""
    weeks = _year_total_weeks(year, method)
    if not 1 <= week <= weeks:
        raise ValueError("week must be in 1..{} for year".format(weeks))

def _check_method(method: str) -> None:
    """Check value of calculation method."""
    methods = ["CDC", "ISO"]
    if method.upper() not in methods:
        raise ValueError(f"method must be '{methods[0]}' or '{methods[1]}'")

def _method_adjustment(method: str) -> int:
    """Return needed adjustment based on first day of week."""
    first_day = ("Mon", "Sun")
    if method.upper() == "CDC":
        return first_day.index("Sun")
    return first_day.index("Mon")

def _year_start(year: int, method: str) -> int:
    """Return ordinal for first day of first week for year."""
    adjustment = _method_adjustment(method)
    mid_weekday = 3 - adjustment  # Sun is 6 .. Mon is 0
    jan1 = date(year, 1, 1)
    jan1_ordinal = jan1.toordinal()
    jan1_weekday = jan1.weekday()
    week1_start_ordinal = jan1_ordinal - jan1_weekday - adjustment
    if jan1_weekday > mid_weekday:
        week1_start_ordinal += 7
    return week1_start_ordinal

def _year_total_weeks(year: int, method: str) -> int:
    """Return number of weeks in year."""
    year_start_ordinal = _year_start(year, method)
    next_year_start_ordinal = _year_start(year + 1, method)
    weeks = (next_year_start_ordinal - year_start_ordinal) // 7
    return weeks
<|endoftext|>