<commit_before>from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.contrib.postgres.search import TrigramDistance
from django.db.models import Q, QuerySet, Value
from django.db.models.functions import Concat, Lower
from django.urls import reverse
from django.utils import timezone
from django.utils.translation import ugettext_lazy as _
from djangorestframework_camel_case.parser import CamelCaseJSONParser
from djangorestframework_camel_case.render import CamelCaseJSONRenderer
from rest_framework import viewsets
from rest_framework.decorators import action
from rest_framework.generics import get_object_or_404
from rest_framework.renderers import BrowsableAPIRenderer
from rest_framework.response import Response
from rest_framework.status import HTTP_400_BAD_REQUEST

from astrobin.models import GearMigrationStrategy, Image
from astrobin_apps_equipment.api.permissions.is_equipment_moderator_or_read_only import IsEquipmentModeratorOrReadOnly
from astrobin_apps_equipment.models import EquipmentItem
from astrobin_apps_equipment.models.equipment_item_group import EquipmentItemKlass
from astrobin_apps_equipment.services.equipment_item_service import EquipmentItemService
from astrobin_apps_notifications.utils import build_notification_url, push_notification
from common.services import AppRedirectionService

class EquipmentItemViewSet(viewsets.ModelViewSet):
    renderer_classes = [BrowsableAPIRenderer, CamelCaseJSONRenderer]
    parser_classes = [CamelCaseJSONParser]
    permission_classes = [IsEquipmentModeratorOrReadOnly]
    http_method_names = ['get', 'post', 'head']

    def get_queryset(self):
        q = self.request.GET.get('q')
        sort = self.request.GET.get('sort')

        manager = self.get_serializer().Meta.model.objects

        if self.request.user.is_authenticated:
            queryset = manager.filter(Q(brand__isnull=False) | Q(created_by=self.request.user))
        else:
            queryset = manager.filter(brand__isnull=False)

        if q:
            queryset = queryset.annotate(
                full_name=Concat('brand__name', Value(' '), 'name'),
                distance=TrigramDistance('full_name', q)
            ).filter(
                Q(distance__lte=.8) | Q(full_name__icontains=q)
            ).order_by(
                'distance'
            )[:10]
        elif sort == "az":
            queryset = queryset.order_by(Lower('brand__name'), Lower('name'))

        return queryset

    @action(
        detail=False,
        methods=['GET'],
        url_path='recently-used'
    )
    def find_recently_used(self, request):
        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        from astrobin_apps_equipment.models import Sensor
        from astrobin_apps_equipment.models import Camera
        from astrobin_apps_equipment.models import Telescope
        from astrobin_apps_equipment.models import Mount
        from astrobin_apps_equipment.models import Filter
        from astrobin_apps_equipment.models import Accessory
        from astrobin_apps_equipment.models import Software

        if manager.model == Sensor:
            return Response("This API does not support sensors", HTTP_400_BAD_REQUEST)

        if request.user.is_authenticated:
            usage_type = request.query_params.get('usage-type')
            recent_items = []
            images: QuerySet[Image] = Image.objects_including_wip.filter(user=request.user).order_by('-uploaded')

            image: Image
            for image in images.iterator():
                if len(recent_items) > 5:
                    break

                property: str
                if manager.model == Camera:
                    if usage_type == 'imaging':
                        property = 'imaging_cameras_2'
                    elif usage_type == 'guiding':
                        property = 'guiding_cameras_2'
                    else:
                        return Response("You need to specify a 'usage_type' with cameras", HTTP_400_BAD_REQUEST)
                elif manager.model == Telescope:
                    if usage_type == 'imaging':
                        property = 'imaging_telescopes_2'
                    elif usage_type == 'guiding':
                        property = 'guiding_telescopes_2'
                    else:
                        return Response("You need to specify a 'usage_type' with cameras", HTTP_400_BAD_REQUEST)
                elif manager.model == Mount:
                    property = 'mounts_2'
                elif manager.model == Filter:
                    property = 'filters_2'
                elif manager.model == Accessory:
                    property = 'accessories_2'
                elif manager.model == Software:
                    property = 'software_2'

                for x in getattr(image, property).all():
                    recent_items.append(x.pk)

            objects = manager.filter(pk__in=recent_items)

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(
        detail=False,
        methods=['GET'],
        url_path='find-similar-in-brand',
    )
    def find_similar_in_brand(self, request):
        brand = request.GET.get('brand')
        q = request.GET.get('q')

        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        if brand and q:
            objects = manager \
                          .annotate(distance=TrigramDistance('name', q)) \
                          .filter(
                Q(brand=int(brand)) &
                Q(Q(distance__lte=.7) | Q(name__icontains=q)) &
                ~Q(name=q)
            ) \
                          .order_by('distance')[:10]

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(
        detail=False,
        methods=['GET'],
        url_path='others-in-brand',
    )
    def others_in_brand(self, request):
        brand = request.GET.get('brand')

        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        if brand:
            objects = manager.filter(brand=int(brand)).order_by('name')

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(detail=True, methods=['POST'])
    def approve(self, request, pk):
        item = get_object_or_404(self.get_serializer().Meta.model.objects, pk=pk)

        if item.reviewed_by is not None:
            return Response("This item was already reviewed", HTTP_400_BAD_REQUEST)

        if item.created_by == request.user:
            return Response("You cannot review an item that you created", HTTP_400_BAD_REQUEST)

        item.reviewed_by = request.user
        item.reviewed_timestamp = timezone.now()
        item.reviewer_decision = 'APPROVED'
        item.reviewer_comment = request.data.get('comment')

        if item.created_by and item.created_by != request.user:
            push_notification(
                [item.created_by],
                request.user,
                'equipment-item-approved',
                {
                    'user': request.user.userprofile.get_display_name(),
                    'user_url': build_notification_url(
                        settings.BASE_URL + reverse('user_page', args=(request.user.username,))
                    ),
                    'item': f'{item.brand.name if item.brand else _("(DIY)")} {item.name}',
                    'item_url': build_notification_url(
                        AppRedirectionService.redirect(
                            f'/equipment/explorer/{EquipmentItemService(item).get_type()}/{item.pk}'
                        )
                    ),
                    'comment': item.reviewer_comment,
                }
            )

        item.save()

        serializer = self.serializer_class(item)
        return Response(serializer.data)

    @action(detail=True, methods=['POST'])
    def reject(self, request, pk):
        from astrobin_apps_equipment.models import Sensor
        from astrobin_apps_equipment.models import Camera
        from astrobin_apps_equipment.models import CameraEditProposal
        from astrobin_apps_equipment.models import Telescope
        from astrobin_apps_equipment.models import Mount
        from astrobin_apps_equipment.models import Filter
        from astrobin_apps_equipment.models import Accessory
        from astrobin_apps_equipment.models import Software

        model = self.get_serializer().Meta.model
        item: EquipmentItem = get_object_or_404(model.objects, pk=pk)

        if item.reviewed_by is not None and item.reviewer_decision == 'APPROVED':
            return Response("This item was already approved", HTTP_400_BAD_REQUEST)

        if item.created_by == request.user:
            return Response("You cannot review an item that you created", HTTP_400_BAD_REQUEST)

        item.reviewed_by = request.user
        item.reviewed_timestamp = timezone.now()
        item.reviewer_decision = 'REJECTED'
        item.reviewer_rejection_reason = request.data.get('reason')
        item.reviewer_comment = request.data.get('comment')

        if item.created_by and item.created_by != request.user:
            push_notification(
                [item.created_by],
                request.user,
                'equipment-item-rejected',
                {
                    'user': request.user.userprofile.get_display_name(),
                    'user_url': build_notification_url(
                        settings.BASE_URL + reverse('user_page', args=(request.user.username,))
                    ),
                    'item': f'{item.brand.name if item.brand else _("(DIY)")} {item.name}',
                    'reject_reason': item.reviewer_rejection_reason,
                    'comment': item.reviewer_comment,
                }
            )

        if item.klass == EquipmentItemKlass.SENSOR:
            Camera.all_objects.filter(sensor=item).update(sensor=None)
            CameraEditProposal.all_objects.filter(sensor=item).update(sensor=None)

        item.delete()

        GearMigrationStrategy.objects.filter(
            migration_flag='MIGRATE',
            migration_content_type=ContentType.objects.get_for_model(model),
            migration_object_id=item.id,
        ).delete()

        if item.brand:
            brand_has_items = False
            for klass in (Sensor, Camera, Telescope, Mount, Filter, Accessory, Software):
                if klass.objects.filter(brand=item.brand).exists():
                    brand_has_items = True
                    break

            if not brand_has_items:
                item.brand.delete()

        serializer = self.serializer_class(item)
        return Response(serializer.data)

    def image_upload(self, request, pk):
        obj = self.get_object()
        serializer = self.serializer_class(obj, data=request.data, partial=True)

        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)

        return Response(serializer.errors, HTTP_400_BAD_REQUEST)

    class Meta:
        abstract = True
<commit_msg>(api/v2) Fix copy/paste mistake in API error message<commit_after>from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.contrib.postgres.search import TrigramDistance
from django.db.models import Q, QuerySet, Value
from django.db.models.functions import Concat, Lower
from django.urls import reverse
from django.utils import timezone
from django.utils.translation import ugettext_lazy as _
from djangorestframework_camel_case.parser import CamelCaseJSONParser
from djangorestframework_camel_case.render import CamelCaseJSONRenderer
from rest_framework import viewsets
from rest_framework.decorators import action
from rest_framework.generics import get_object_or_404
from rest_framework.renderers import BrowsableAPIRenderer
from rest_framework.response import Response
from rest_framework.status import HTTP_400_BAD_REQUEST

from astrobin.models import GearMigrationStrategy, Image
from astrobin_apps_equipment.api.permissions.is_equipment_moderator_or_read_only import IsEquipmentModeratorOrReadOnly
from astrobin_apps_equipment.models import EquipmentItem
from astrobin_apps_equipment.models.equipment_item_group import EquipmentItemKlass
from astrobin_apps_equipment.services.equipment_item_service import EquipmentItemService
from astrobin_apps_notifications.utils import build_notification_url, push_notification
from common.services import AppRedirectionService

class EquipmentItemViewSet(viewsets.ModelViewSet):
    renderer_classes = [BrowsableAPIRenderer, CamelCaseJSONRenderer]
    parser_classes = [CamelCaseJSONParser]
    permission_classes = [IsEquipmentModeratorOrReadOnly]
    http_method_names = ['get', 'post', 'head']

    def get_queryset(self):
        q = self.request.GET.get('q')
        sort = self.request.GET.get('sort')

        manager = self.get_serializer().Meta.model.objects

        if self.request.user.is_authenticated:
            queryset = manager.filter(Q(brand__isnull=False) | Q(created_by=self.request.user))
        else:
            queryset = manager.filter(brand__isnull=False)

        if q:
            queryset = queryset.annotate(
                full_name=Concat('brand__name', Value(' '), 'name'),
                distance=TrigramDistance('full_name', q)
            ).filter(
                Q(distance__lte=.8) | Q(full_name__icontains=q)
            ).order_by(
                'distance'
            )[:10]
        elif sort == "az":
            queryset = queryset.order_by(Lower('brand__name'), Lower('name'))

        return queryset

    @action(
        detail=False,
        methods=['GET'],
        url_path='recently-used'
    )
    def find_recently_used(self, request):
        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        from astrobin_apps_equipment.models import Sensor
        from astrobin_apps_equipment.models import Camera
        from astrobin_apps_equipment.models import Telescope
        from astrobin_apps_equipment.models import Mount
        from astrobin_apps_equipment.models import Filter
        from astrobin_apps_equipment.models import Accessory
        from astrobin_apps_equipment.models import Software

        if manager.model == Sensor:
            return Response("This API does not support sensors", HTTP_400_BAD_REQUEST)

        if request.user.is_authenticated:
            usage_type = request.query_params.get('usage-type')
            recent_items = []
            images: QuerySet[Image] = Image.objects_including_wip.filter(user=request.user).order_by('-uploaded')

            image: Image
            for image in images.iterator():
                if len(recent_items) > 5:
                    break

                property: str
                if manager.model == Camera:
                    if usage_type == 'imaging':
                        property = 'imaging_cameras_2'
                    elif usage_type == 'guiding':
                        property = 'guiding_cameras_2'
                    else:
                        return Response("You need to specify a 'usage_type' with cameras", HTTP_400_BAD_REQUEST)
                elif manager.model == Telescope:
                    if usage_type == 'imaging':
                        property = 'imaging_telescopes_2'
                    elif usage_type == 'guiding':
                        property = 'guiding_telescopes_2'
                    else:
                        return Response("You need to specify a 'usage_type' with telescopes", HTTP_400_BAD_REQUEST)
                elif manager.model == Mount:
                    property = 'mounts_2'
                elif manager.model == Filter:
                    property = 'filters_2'
                elif manager.model == Accessory:
                    property = 'accessories_2'
                elif manager.model == Software:
                    property = 'software_2'

                for x in getattr(image, property).all():
                    recent_items.append(x.pk)

            objects = manager.filter(pk__in=recent_items)

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(
        detail=False,
        methods=['GET'],
        url_path='find-similar-in-brand',
    )
    def find_similar_in_brand(self, request):
        brand = request.GET.get('brand')
        q = request.GET.get('q')

        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        if brand and q:
            objects = manager \
                          .annotate(distance=TrigramDistance('name', q)) \
                          .filter(
                Q(brand=int(brand)) &
                Q(Q(distance__lte=.7) | Q(name__icontains=q)) &
                ~Q(name=q)
            ) \
                          .order_by('distance')[:10]

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(
        detail=False,
        methods=['GET'],
        url_path='others-in-brand',
    )
    def others_in_brand(self, request):
        brand = request.GET.get('brand')

        manager = self.get_serializer().Meta.model.objects
        objects = manager.none()

        if brand:
            objects = manager.filter(brand=int(brand)).order_by('name')

        serializer = self.serializer_class(objects, many=True)
        return Response(serializer.data)

    @action(detail=True, methods=['POST'])
    def approve(self, request, pk):
        item = get_object_or_404(self.get_serializer().Meta.model.objects, pk=pk)

        if item.reviewed_by is not None:
            return Response("This item was already reviewed", HTTP_400_BAD_REQUEST)

        if item.created_by == request.user:
            return Response("You cannot review an item that you created", HTTP_400_BAD_REQUEST)

        item.reviewed_by = request.user
        item.reviewed_timestamp = timezone.now()
        item.reviewer_decision = 'APPROVED'
        item.reviewer_comment = request.data.get('comment')

        if item.created_by and item.created_by != request.user:
            push_notification(
                [item.created_by],
                request.user,
                'equipment-item-approved',
                {
                    'user': request.user.userprofile.get_display_name(),
                    'user_url': build_notification_url(
                        settings.BASE_URL + reverse('user_page', args=(request.user.username,))
                    ),
                    'item': f'{item.brand.name if item.brand else _("(DIY)")} {item.name}',
                    'item_url': build_notification_url(
                        AppRedirectionService.redirect(
                            f'/equipment/explorer/{EquipmentItemService(item).get_type()}/{item.pk}'
                        )
                    ),
                    'comment': item.reviewer_comment,
                }
            )

        item.save()

        serializer = self.serializer_class(item)
        return Response(serializer.data)

    @action(detail=True, methods=['POST'])
    def reject(self, request, pk):
        from astrobin_apps_equipment.models import Sensor
        from astrobin_apps_equipment.models import Camera
        from astrobin_apps_equipment.models import CameraEditProposal
        from astrobin_apps_equipment.models import Telescope
        from astrobin_apps_equipment.models import Mount
        from astrobin_apps_equipment.models import Filter
        from astrobin_apps_equipment.models import Accessory
        from astrobin_apps_equipment.models import Software

        model = self.get_serializer().Meta.model
        item: EquipmentItem = get_object_or_404(model.objects, pk=pk)

        if item.reviewed_by is not None and item.reviewer_decision == 'APPROVED':
            return Response("This item was already approved", HTTP_400_BAD_REQUEST)

        if item.created_by == request.user:
            return Response("You cannot review an item that you created", HTTP_400_BAD_REQUEST)

        item.reviewed_by = request.user
        item.reviewed_timestamp = timezone.now()
        item.reviewer_decision = 'REJECTED'
        item.reviewer_rejection_reason = request.data.get('reason')
        item.reviewer_comment = request.data.get('comment')

        if item.created_by and item.created_by != request.user:
            push_notification(
                [item.created_by],
                request.user,
                'equipment-item-rejected',
                {
                    'user': request.user.userprofile.get_display_name(),
                    'user_url': build_notification_url(
                        settings.BASE_URL + reverse('user_page', args=(request.user.username,))
                    ),
                    'item': f'{item.brand.name if item.brand else _("(DIY)")} {item.name}',
                    'reject_reason': item.reviewer_rejection_reason,
                    'comment': item.reviewer_comment,
                }
            )

        if item.klass == EquipmentItemKlass.SENSOR:
            Camera.all_objects.filter(sensor=item).update(sensor=None)
            CameraEditProposal.all_objects.filter(sensor=item).update(sensor=None)

        item.delete()

        GearMigrationStrategy.objects.filter(
            migration_flag='MIGRATE',
            migration_content_type=ContentType.objects.get_for_model(model),
            migration_object_id=item.id,
        ).delete()

        if item.brand:
            brand_has_items = False
            for klass in (Sensor, Camera, Telescope, Mount, Filter, Accessory, Software):
                if klass.objects.filter(brand=item.brand).exists():
                    brand_has_items = True
                    break

            if not brand_has_items:
                item.brand.delete()

        serializer = self.serializer_class(item)
        return Response(serializer.data)

    def image_upload(self, request, pk):
        obj = self.get_object()
        serializer = self.serializer_class(obj, data=request.data, partial=True)

        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)

        return Response(serializer.errors, HTTP_400_BAD_REQUEST)

    class Meta:
        abstract = True
<|endoftext|>


<commit_before># This code is written for python 2.x
#-*- coding: utf-8 -*-
import urllib2
import json
import string
import re
from bs4 import BeautifulSoup
import logging
import time

FILE_PATH = "./boxofficemojo.com/movie_data.txt"
LOG_PATH = "./boxofficemojo.com/scraping.log"

logging.basicConfig(filename=LOG_PATH,level=logging.DEBUG)

Keys = ["Name", "URL", "Genre","Runtime", "Rating", "MovieRanking"
    , "PercentageofTotalGross", "WidestRelease", "CloseDate", "InRelease", "TotalGross"
    , "Distributor", "Budget", "Domestic_Gross", "Domestic_Percentage"
    , "Foreign_Gross", "Foreign_Percentage", "Worldwide_Gross", "OpeningWeekend"
    , "Countryclicktoviewweekendbreakdown", "Dist", "ReleaseDate"
    , "OpeningWknd", "ofTotal", "TotalGross", "AsOf"]

def add_empty_data(arrData, count):
    for i in range(0,count):
        arrData.append(" ")
    return arrData

def remove_special_chars(dictData):

    newDict= {}
    for key in dictData:
        new_key= re.sub(r'\W+', '', key)
        newDict[new_key] = dictData[key]
    return newDict

def save_to_json(filePath, dictData, countriesData=None):

    dictData = remove_special_chars(dictData)
    countriesData = remove_special_chars(countriesData)

    if countriesData:
        merged = dict(dictData)
        merged.update(countriesData)
        dictData = merged

    with open(filePath, "a") as outfile:
        json.dump(dictData, outfile, ensure_ascii=False)

def write_header(filePath):
    # Write a header
    text_file = open(filePath, "ab")

    for header in Keys:
        text_file.write((header + u"|").encode('utf-8'))
    text_file.write("\n".encode('utf-8'))
    text_file.close()

def save_to_file(filePath, dictData, countriesData=None):

    dictData = remove_special_chars(dictData)

    if countriesData:
        countriesData = remove_special_chars(countriesData)

    if countriesData:
        merged = dict(dictData)
        merged.update(countriesData)
        dictData = merged

    Arranged= []
    add_empty_data(Arranged, 50)

    text_file = open(filePath, "ab")
    for key, value in dictData.items():
        for i ,k in enumerate(Keys):
            if key == k:
                Arranged[i]= value

    for data in Arranged:
        text_file.write((data + u"|").encode('utf-8'))

    text_file.write("\n".encode('utf-8'))
    text_file.close()

def get_total_lifetime_grosses(link, arrData):

    url = "http://www.boxofficemojo.com"+ link
    page = urllib2.urlopen(url)
    soup = BeautifulSoup(page, "lxml")

    # Assume that domestic info is from USA
    arrData['Countryclicktoviewweekendbreakdown']= "USA"
    #print(main_tbl)
    tables = soup.find_all('table', attrs={'border': '0' , 'cellspacing':'0', 'cellpadding':'0' , 'width':'100%'})

    #print( len(tables))
    #td_count = 9
    if len(tables) == 4:
        #print(tables[3]) # Total lifetime grosses
        mp_boxes= tables[3].find_all("div", {"class", "mp_box_tab"})
        a= len(mp_boxes)
        for box in mp_boxes:
            if(box.text == "Total Lifetime Grosses"):
                div_content= box.findNext('div')
                trs = div_content.find_all('tr')
                for tr in trs:
                    tds = tr.find_all('td')
                    if len(tds) == 3:
                        if tds[0].text.strip() == "Domestic:":
                            arrData["Total Gross"] = tds[1].text.strip()
                            arrData["% ofTotal"] = tds[2].text.strip()
                        arrData[tds[0].text.strip()+"_Gross"] = tds[1].text.strip()
                        arrData[tds[0].text.strip()+"_Percentage"] = tds[2].text.strip()

            if(box.text == "Domestic Summary"):
                div_content = box.findNext('div')
                DS_tables = div_content.find_all('table', attrs = { 'border': '0' , 'cellspacing':'0', 'cellpadding':'0'})
                for DS_table in DS_tables:
                    DS_trs = DS_table.find_all('tr')
                    for DS_tr in DS_trs:
                        DS_tr_title = DS_tr.td.text.strip()
                        if(DS_tr_title == "Opening\xa0Weekend:") or (DS_tr_title == "Opening Weekend:"):
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData["Opening Weekend"] = DS_tr_content.text.strip()
                                arrData["OpeningWknd"] = DS_tr_content.text.strip()

                        elif "(#" in DS_tr_title:
                            arrData['Movie Ranking'] = DS_tr_title

                        elif "%\xa0of\xa0Total\xa0Gross" in DS_tr_title or "% of Total Gross" in DS_tr_title:
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Percentage of Total Gross'] = DS_tr_content.text.strip()

                        elif DS_tr_title == "Widest\xa0Release:" or DS_tr_title == "Widest Release:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Widest Release'] = DS_tr_content.text.strip() # 14.

                        elif DS_tr_title == "Close\xa0Date:" or DS_tr_title == "Close Date:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Close Date'] = DS_tr_content.text.strip() # 15.

                        elif DS_tr_title == "In\xa0Release:" or DS_tr_title == "In Release:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['In Release'] = DS_tr_content.text.strip() # 15.

            if(box.text == "The Players"):
                #print(box.findNext('div'))
                pass

    return arrData

def get_movie_foreign(link, arrData):

    try:
        eachCountry = {}
        ColumnHeaders= []
        url = "http://www.boxofficemojo.com"+ link + "&page=intl"
        page = urllib2.urlopen(url)
        soup = BeautifulSoup(page, "lxml")

        contents = soup.find('table', attrs={'border': '3' , 'cellspacing':'0', 'cellpadding':'5', 'align':'center', 'style':'margin-top: 5px;'})
        if len(contents) == 1:
            #print(contents)
            intl_table = contents.tr.td.table
            if intl_table:
                trs = intl_table.find_all("tr")
                if len(trs) == 3:
                    #print ("no data")
                    temp= 0
                else:
                    for row,tr in enumerate(trs):
                        if row == 0:
                            tds= tr.find_all("td") # get each header's text
                            for td in tds:
                                header= td.text.strip()
                                if "/" in header:
                                    divided_header = header.split('/')
                                    ColumnHeaders.append(divided_header[0])
                                    ColumnHeaders.append(divided_header[1])
                                else:
                                    ColumnHeaders.append(td.text.strip())
                        if(row < 3): # don't save unncessary data
                            continue
                        tds= tr.find_all("td")
                        for column, td in enumerate(tds):
                            # 11. Country, 12.Dist, 13. Release Date, 14.OW, 15.% of Total, 16.Total gross, 17. as of
                            eachCountry[ColumnHeaders[column]] = td.text.strip()
                        save_to_file(FILE_PATH, arrData, eachCountry)
                        #save_to_json(FILE_PATH, arrData, eachCountry)
                        eachCountry.clear()
        return arrData
    except Exception as e:
        logging.exception(e)
        return arrData

def get_movie_detail(movies_list, link, arrData):

    if link not in movies_list:
        movies_list.append(link)

        url = "http://www.boxofficemojo.com"+ link # 1. URL
        page = urllib2.urlopen(url)
        soup = BeautifulSoup(page, "lxml")
        contents= soup.find('table', attrs={'border': '0' , 'cellspacing':'1', 'cellpadding':'4' , 'bgcolor':'#dcdcdc', 'width':'95%'})
        tabledata = contents.find_all("td")

        name_table = soup.find('table', attrs={'border': '0' , 'cellspacing':'0', 'cellpadding':'0' , 'width':'100%', 'style':'padding-top: 5px;'})
        name = name_table.font.b.getText() # 0. Name

        # 2. Distributor, 3. Release Date, 4. Genre, 5. Runtime, 6. Rating, 7. Budget, 8. TotalGross
        arrData['Name'] = name
        arrData['URL'] = url
        if len(tabledata) == 6:
            Distributor = tabledata[0].b.getText()
            ReleaseDate = tabledata[1].b.getText()
            Genre = tabledata[2].b.getText()
            Runtime = tabledata[3].b.getText()
            Rating = tabledata[4].b.getText()
            Budget = tabledata[5].b.getText()

            arrData['Distributor'] = Distributor
            arrData['ReleaseDate'] = ReleaseDate
            arrData['Genre'] = Genre
            arrData['Runtime'] = Runtime
            arrData['Rating'] = Rating
            arrData['Budget'] = Budget

            #arrData.extend([name , url , Distributor, ReleaseDate,Genre ,Runtime , Rating,Budget])
            #add_empty_data(arrData, 1) # match gap for missing column
        elif len(tabledata) == 7:
            TotalGross = tabledata[0].b.getText()
            Distributor = tabledata[1].b.getText()
            ReleaseDate = tabledata[2].b.getText()
            Genre = tabledata[3].b.getText()
            Runtime = tabledata[4].b.getText()
            Rating = tabledata[5].b.getText()
            Budget = tabledata[6].b.getText()

            arrData['TotalGross'] = TotalGross
            arrData['Distributor'] = Distributor
            arrData['ReleaseDate'] = ReleaseDate
            arrData['Genre'] = Genre
            arrData['Runtime'] = Runtime
            arrData['Rating'] = Rating
            arrData['Budget'] = Budget
            #arrData.extend([ name , url , Distributor, ReleaseDate,Genre ,Runtime , Rating,Budget ,TotalGross])

            #print (result)

        #print contents2[0]
    return arrData

def get_all_movies():
    # Alphabet loop for how movies are indexed including
    # movies that start with a special character or number
    index = ["NUM"] + list(string.ascii_uppercase)

    # List of movie urls
    movies_list = []

    # dict data
    arrData = {}

    startTime = time.time()
    lapTime= 0.0

    write_header(FILE_PATH)
    logging.debug("running...start at : " + str(time.time()))
    # Loop through the pages for each letter
    for letter in index:

        url = ("http://www.boxofficemojo.com/movies/alphabetical.htm?letter=" + letter)
        page1 = urllib2.urlopen(url)
        soup1 = BeautifulSoup(page1, "lxml")
        navi = soup1.find('div', attrs={"class" : "alpha-nav-holder"})
        bs= navi.font.find_all('b')
        count_bs= len(bs)
        logging.debug("pages count : " + str(count_bs))

        if letter == "NUM":
            count_bs = 1

        # Loop through the pages within each letter
        for num in range(1, count_bs+1):
            logging.debug("begin to scrap letter : " + letter + ", page : " + str(num))

            url = ("http://www.boxofficemojo.com/movies/alphabetical.htm?"
                   "letter=" + letter + "&page=" + str(num))
            try:
                page = urllib2.urlopen(url)
                soup = BeautifulSoup(page, "lxml")
                rows = soup.find(id="body").find("table").find("table").find_all(
                    "table")[1].find_all("tr")

                # skip index row
                if len(rows) > 1:
                    counter = 1
                    for row in rows:

                        trackingStartTime= time.time()
                        # skip index row
                        if counter > 1:
                            link = row.td.font.a['href']
                            arrData = get_movie_detail(movies_list, link, arrData)
                            arrData = get_movie_foreign(link, arrData)
                            arrData = get_total_lifetime_grosses(link, arrData)
                            save_to_file(FILE_PATH, arrData)
                            arrData.clear()

                            lapTime= time.time() - trackingStartTime
                            logging.debug("each movie's lapTime : " + str(lapTime))
                        counter += 1
            except Exception as e:
                logging.exception(e)

    TotalElaspedTime= (time.time() - startTime)
    logging.debug('done.' + str(TotalElaspedTime))

get_all_movies()<commit_msg>Make it unicode<commit_after># This code is written for python 2.x
#-*- coding: utf-8 -*-
import urllib2
import json
import string
import re
from bs4 import BeautifulSoup
import logging
import time

FILE_PATH = "./boxofficemojo.com/movie_data.txt"
LOG_PATH = "./boxofficemojo.com/scraping.log"

logging.basicConfig(filename=LOG_PATH,level=logging.DEBUG)

Keys = ["Name", "URL", "Genre","Runtime", "Rating", "MovieRanking"
    , "PercentageofTotalGross", "WidestRelease", "CloseDate", "InRelease", "TotalGross"
    , "Distributor", "Budget", "Domestic_Gross", "Domestic_Percentage"
    , "Foreign_Gross", "Foreign_Percentage", "Worldwide_Gross", "OpeningWeekend"
    , "Countryclicktoviewweekendbreakdown", "Dist", "ReleaseDate"
    , "OpeningWknd", "ofTotal", "TotalGross", "AsOf"]

def add_empty_data(arrData, count):
    for i in range(0,count):
        arrData.append(" ")
    return arrData

def remove_special_chars(dictData):

    newDict= {}
    for key in dictData:
        new_key= re.sub(r'\W+', '', key)
        newDict[new_key] = dictData[key]
    return newDict

def save_to_json(filePath, dictData, countriesData=None):

    dictData = remove_special_chars(dictData)
    countriesData = remove_special_chars(countriesData)

    if countriesData:
        merged = dict(dictData)
        merged.update(countriesData)
        dictData = merged

    with open(filePath, "a") as outfile:
        json.dump(dictData, outfile, ensure_ascii=False)

def write_header(filePath):
    # Write a header
    text_file = open(filePath, "ab")

    for header in Keys:
        text_file.write((header + u"|").encode('utf-8'))
    text_file.write("\n".encode('utf-8'))
    text_file.close()

def save_to_file(filePath, dictData, countriesData=None):

    dictData = remove_special_chars(dictData)

    if countriesData:
        countriesData = remove_special_chars(countriesData)

    if countriesData:
        merged = dict(dictData)
        merged.update(countriesData)
        dictData = merged

    Arranged= []
    add_empty_data(Arranged, 50)

    text_file = open(filePath, "ab")
    for key, value in dictData.items():
        for i ,k in enumerate(Keys):
            if key == k:
                Arranged[i]= value

    for data in Arranged:
        text_file.write((data + u"|").encode('utf-8'))

    text_file.write("\n".encode('utf-8'))
    text_file.close()

def get_total_lifetime_grosses(link, arrData):

    url = "http://www.boxofficemojo.com"+ link
    page = urllib2.urlopen(url)
    soup = BeautifulSoup(page, "lxml")

    # Assume that domestic info is from USA
    arrData['Countryclicktoviewweekendbreakdown']= "USA"
    #print(main_tbl)
    tables = soup.find_all('table', attrs={'border': '0' , 'cellspacing':'0', 'cellpadding':'0' , 'width':'100%'})

    #print( len(tables))
    #td_count = 9
    if len(tables) == 4:
        #print(tables[3]) # Total lifetime grosses
        mp_boxes= tables[3].find_all("div", {"class", "mp_box_tab"})
        a= len(mp_boxes)
        for box in mp_boxes:
            if(box.text == "Total Lifetime Grosses"):
                div_content= box.findNext('div')
                trs = div_content.find_all('tr')
                for tr in trs:
                    tds = tr.find_all('td')
                    if len(tds) == 3:
                        if tds[0].text.strip() == "Domestic:":
                            arrData["Total Gross"] = tds[1].text.strip()
                            arrData["% ofTotal"] = tds[2].text.strip()
                        arrData[tds[0].text.strip()+"_Gross"] = tds[1].text.strip()
                        arrData[tds[0].text.strip()+"_Percentage"] = tds[2].text.strip()

            if(box.text == "Domestic Summary"):
                div_content = box.findNext('div')
                DS_tables = div_content.find_all('table', attrs = { 'border': '0' , 'cellspacing':'0', 'cellpadding':'0'})
                for DS_table in DS_tables:
                    DS_trs = DS_table.find_all('tr')
                    for DS_tr in DS_trs:
                        DS_tr_title = DS_tr.td.text.strip()
                        if(DS_tr_title == "Opening\xa0Weekend:".encode('utf-8')) or (DS_tr_title == "Opening Weekend:"):
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData["Opening Weekend"] = DS_tr_content.text.strip()
                                arrData["OpeningWknd"] = DS_tr_content.text.strip()

                        elif "(#" in DS_tr_title:
                            arrData['Movie Ranking'] = DS_tr_title

                        elif "%\xa0of\xa0Total\xa0Gross".encode('utf-8') in DS_tr_title or "% of Total Gross" in DS_tr_title:
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Percentage of Total Gross'] = DS_tr_content.text.strip()

                        elif DS_tr_title == "Widest\xa0Release:".encode('utf-8') or DS_tr_title == "Widest Release:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Widest Release'] = DS_tr_content.text.strip() # 14.

                        elif DS_tr_title == "Close\xa0Date:".encode('utf-8') or DS_tr_title == "Close Date:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['Close Date'] = DS_tr_content.text.strip() # 15.

                        elif DS_tr_title == "In\xa0Release:".encode('utf-8') or DS_tr_title == "In Release:":
                            DS_tr_content = DS_tr.td.findNext('td')
                            if DS_tr_content:
                                arrData['In Release'] = DS_tr_content.text.strip() # 15.

            if(box.text == "The Players"):
                #print(box.findNext('div'))
                pass

    return arrData

def get_movie_foreign(link, arrData):

    try:
        eachCountry = {}
        ColumnHeaders= []
        url = "http://www.boxofficemojo.com"+ link + "&page=intl"
        page = urllib2.urlopen(url)
        soup = BeautifulSoup(page, "lxml")

        contents = soup.find('table', attrs={'border': '3' , 'cellspacing':'0', 'cellpadding':'5', 'align':'center', 'style':'margin-top: 5px;'})
        if len(contents) == 1:
            #print(contents)
            intl_table = contents.tr.td.table
            if intl_table:
                trs = intl_table.find_all("tr")
                if len(trs) == 3:
                    #print ("no data")
                    temp= 0
                else:
                    for row,tr in enumerate(trs):
                        if row == 0:
                            tds= tr.find_all("td") # get each header's text
                            for td in tds:
                                header= td.text.strip()
                                if "/" in header:
                                    divided_header = header.split('/')
                                    ColumnHeaders.append(divided_header[0])
                                    ColumnHeaders.append(divided_header[1])
                                else:
                                    ColumnHeaders.append(td.text.strip())
                        if(row < 3): # don't save unncessary data
                            continue
                        tds= tr.find_all("td")
                        for column, td in enumerate(tds):
                            # 11. Country, 12.Dist, 13. Release Date, 14.OW, 15.% of Total, 16.Total gross, 17. as of
                            eachCountry[ColumnHeaders[column]] = td.text.strip()
                        save_to_file(FILE_PATH, arrData, eachCountry)
                        #save_to_json(FILE_PATH, arrData, eachCountry)
                        eachCountry.clear()
        return arrData
    except Exception as e:
        logging.exception(e)
        return arrData

def get_movie_detail(movies_list, link, arrData):

    if link not in movies_list:
        movies_list.append(link)

        url = "http://www.boxofficemojo.com"+ link # 1. URL
        page = urllib2.urlopen(url)
        soup = BeautifulSoup(page, "lxml")
        contents= soup.find('table', attrs={'border': '0' , 'cellspacing':'1', 'cellpadding':'4' , 'bgcolor':'#dcdcdc', 'width':'95%'})
        tabledata = contents.find_all("td")

        name_table = soup.find('table', attrs={'border': '0' , 'cellspacing':'0', 'cellpadding':'0' , 'width':'100%', 'style':'padding-top: 5px;'})
        name = name_table.font.b.getText() # 0. Name

        # 2. Distributor, 3. Release Date, 4. Genre, 5. Runtime, 6. Rating, 7. Budget, 8. TotalGross
        arrData['Name'] = name
        arrData['URL'] = url
        if len(tabledata) == 6:
            Distributor = tabledata[0].b.getText()
            ReleaseDate = tabledata[1].b.getText()
            Genre = tabledata[2].b.getText()
            Runtime = tabledata[3].b.getText()
            Rating = tabledata[4].b.getText()
            Budget = tabledata[5].b.getText()

            arrData['Distributor'] = Distributor
            arrData['ReleaseDate'] = ReleaseDate
            arrData['Genre'] = Genre
            arrData['Runtime'] = Runtime
            arrData['Rating'] = Rating
            arrData['Budget'] = Budget

            #arrData.extend([name , url , Distributor, ReleaseDate,Genre ,Runtime , Rating,Budget])
            #add_empty_data(arrData, 1) # match gap for missing column
        elif len(tabledata) == 7:
            TotalGross = tabledata[0].b.getText()
            Distributor = tabledata[1].b.getText()
            ReleaseDate = tabledata[2].b.getText()
            Genre = tabledata[3].b.getText()
            Runtime = tabledata[4].b.getText()
            Rating = tabledata[5].b.getText()
            Budget = tabledata[6].b.getText()

            arrData['TotalGross'] = TotalGross
            arrData['Distributor'] = Distributor
            arrData['ReleaseDate'] = ReleaseDate
            arrData['Genre'] = Genre
            arrData['Runtime'] = Runtime
            arrData['Rating'] = Rating
            arrData['Budget'] = Budget
            #arrData.extend([ name , url , Distributor, ReleaseDate,Genre ,Runtime , Rating,Budget ,TotalGross])

            #print (result)

        #print contents2[0]
    return arrData

def get_all_movies():
    # Alphabet loop for how movies are indexed including
    # movies that start with a special character or number
    index = ["NUM"] + list(string.ascii_uppercase)

    # List of movie urls
    movies_list = []

    # dict data
    arrData = {}

    startTime = time.time()
    lapTime= 0.0

    write_header(FILE_PATH)
    logging.debug("running...start at : " + str(time.time()))
    # Loop through the pages for each letter
    for letter in index:

        url = ("http://www.boxofficemojo.com/movies/alphabetical.htm?letter=" + letter)
        page1 = urllib2.urlopen(url)
        soup1 = BeautifulSoup(page1, "lxml")
        navi = soup1.find('div', attrs={"class" : "alpha-nav-holder"})
        bs= navi.font.find_all('b')
        count_bs= len(bs)
        logging.debug("pages count : " + str(count_bs))

        if letter == "NUM":
            count_bs = 1

        # Loop through the pages within each letter
        for num in range(1, count_bs+1):
            logging.debug("begin to scrap letter : " + letter + ", page : " + str(num))

            url = ("http://www.boxofficemojo.com/movies/alphabetical.htm?"
                   "letter=" + letter + "&page=" + str(num))
            try:
                page = urllib2.urlopen(url)
                soup = BeautifulSoup(page, "lxml")
                rows = soup.find(id="body").find("table").find("table").find_all(
                    "table")[1].find_all("tr")

                # skip index row
                if len(rows) > 1:
                    counter = 1
                    for row in rows:

                        trackingStartTime= time.time()
                        # skip index row
                        if counter > 1:
                            link = row.td.font.a['href']
                            arrData = get_movie_detail(movies_list, link, arrData)
                            arrData = get_movie_foreign(link, arrData)
                            arrData = get_total_lifetime_grosses(link, arrData)
                            save_to_file(FILE_PATH, arrData)
                            arrData.clear()

                            lapTime= time.time() - trackingStartTime
                            logging.debug("each movie's lapTime : " + str(lapTime))
                        counter += 1
            except Exception as e:
                logging.exception(e)

    TotalElaspedTime= (time.time() - startTime)
    logging.debug('done.' + str(TotalElaspedTime))

get_all_movies()<|endoftext|>


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

__author__ = "Sebastián J. Seba"
__version__ = "1.0"
__license__ = """This program is free software: you can redistribute it and/or
                 modify it under the terms of the GNU 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 General Public License for more details.
                 You should have received a copy of the GNU General Public License
                 along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 """

from optparse import OptionParser
from time import sleep
import serial

def main():
    parser = OptionParser(usage="Usage: %prog [-p PORT_NUMBER]" \
                          " [-c CALL_DELAY]\n                 " \
                          "  [-u HANGUP_DELAY] [-r REDIAL_DELAY]\n" \
                          "                   [-a] archivo_llamadas",
                          version=__version__)
    parser.add_option("-p", "--port", action="store", type="int",
                      dest="port_number",
                      help="Numero de puerto COM a utilizar [Obligatorio]")
    parser.add_option("-c", "--call-delay", action="store", type="int",
                      dest="call_delay", default=30,
                      help="Tiempo (en segundos) a esperar luego de realizar " \
                      "una llamada [Default: %default]")
    parser.add_option("-u", "--hangup-delay", action="store", type="int",
                      dest="hangup_delay", default=2,
                      help="Tiempo (en segundos) a esperar luego de cortar " \
                      "una llamada [Default: %default]")
    parser.add_option("-r", "--redial-delay", action="store", type="int",
                      dest="redial_delay", default=2,
                      help="Tiempo (en segundos) a esperar entre llamadas" \
                      "[Default: %default]")
    parser.add_option("-a", "--autoclose", action="store_true", default=False,
                      dest="autoclose", help="Cerrar programa al finalizar. " \
                      "Al utilizar esta opcion no sera necesario presionar " \
                      "la tecla ENTER para salir.")
    (options, args) = parser.parse_args()

    if len(args) == 0:
        parser.print_usage()
        print "Pruebe '{0} --help' para mas" \
               "informacion.".format(parser.get_prog_name())
    elif options.port_number is None:
        print 'Debe proporcionar un numero de puerto COM.'
    else:
        with open(args[0], 'r') as fcalls:
            calls = fcalls.readlines()

        if calls:
            try:
                print u'Conectando a módem en {0}...'.format(serial.device(
                                                             options.port_number - 1))
                modem = serial.Serial(options.port_number - 1, timeout=1)

                total = len(calls)

                for i, call in enumerate(calls, 1):
                    print '{0}/{1}] Marcando: {2}'.format(str(i), total, call)
                    modem.write('ATD' + call.strip() + '\r')
                    sleep(options.call_delay)
                    print modem.read(modem.inWaiting())

                    print 'Colgando...'
                    modem.write('ATH\r')
                    sleep(options.hangup_delay)
                    print modem.read(modem.inWaiting())
                    sleep(options.redial_delay)

                print u'Cerrando conexión...'
                modem.close()
                print 'Terminado. Presione ENTER para salir.'
                raw_input()
            except serial.SerialException:
                print 'Error al conectarse al puerto' \
                      '{0} ({1})'.format(options.port_number - 1,
                                         serial.device(options.port_number - 1))
            except Exception:
                pass
        else:
            print 'El archivo de llamadas no contiene ninguna llamada.'

if __name__ == '__main__':
    main()
<commit_msg>Mejorados mensajes de información de tareas.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

__author__ = "Sebastián J. Seba"
__version__ = "1.0"
__license__ = """This program is free software: you can redistribute it and/or
                 modify it under the terms of the GNU 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 General Public License for more details.
                 You should have received a copy of the GNU General Public License
                 along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 """

from optparse import OptionParser
from time import sleep
import serial

def main():
    parser = OptionParser(usage="Usage: %prog [-p PORT_NUMBER]" \
                          " [-c CALL_DELAY]\n                 " \
                          "  [-u HANGUP_DELAY] [-r REDIAL_DELAY]\n" \
                          "                   [-a] archivo_llamadas",
                          version=__version__)
    parser.add_option("-p", "--port", action="store", type="int",
                      dest="port_number",
                      help="Numero de puerto COM a utilizar [Obligatorio]")
    parser.add_option("-c", "--call-delay", action="store", type="int",
                      dest="call_delay", default=30,
                      help="Tiempo (en segundos) a esperar luego de realizar " \
                      "una llamada [Default: %default]")
    parser.add_option("-u", "--hangup-delay", action="store", type="int",
                      dest="hangup_delay", default=2,
                      help="Tiempo (en segundos) a esperar luego de cortar " \
                      "una llamada [Default: %default]")
    parser.add_option("-r", "--redial-delay", action="store", type="int",
                      dest="redial_delay", default=2,
                      help="Tiempo (en segundos) a esperar entre llamadas" \
                      "[Default: %default]")
    parser.add_option("-a", "--autoclose", action="store_true", default=False,
                      dest="autoclose", help="Cerrar programa al finalizar. " \
                      "Al utilizar esta opcion no sera necesario presionar " \
                      "la tecla ENTER para salir.")
    (options, args) = parser.parse_args()

    if len(args) == 0:
        parser.print_usage()
        print "Pruebe '{0} --help' para mas " \
               "informacion.".format(parser.get_prog_name())
    elif options.port_number is None:
        print 'Debe proporcionar un numero de puerto COM.'
    else:
        with open(args[0], 'r') as fcalls:
            calls = fcalls.readlines()

        if calls:
            try:
                print u'Conectando a modem en {0}...'.format(serial.device(
                                                             options.port_number - 1))
                modem = serial.Serial(options.port_number - 1, timeout=1)

                total = len(calls)

                for i, call in enumerate(calls, 1):
                    print '{0}/{1}] Marcando: {2}'.format(str(i), total, call)
                    modem.write('ATD' + call.strip() + '\r')
                    sleep(options.call_delay)
                    print modem.read(modem.inWaiting())

                    print 'Colgando...'
                    modem.write('ATH\r')
                    sleep(options.hangup_delay)
                    print modem.read(modem.inWaiting())
                    sleep(options.redial_delay)

                print u'Cerrando conexión con modem...'
                modem.close()

                if not options.autoclose:
                    print 'Terminado. Presione ENTER para salir.'
                    raw_input()
            except serial.SerialException:
                print 'Error al conectarse al puerto' \
                      '{0} ({1})'.format(options.port_number - 1,
                                         serial.device(options.port_number - 1))
            except Exception:
                pass
        else:
            print 'El archivo de llamadas no contiene ninguna llamada.'

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>#!/usr/bin/python

# Copyright 2015 Huawei Devices USA Inc. All rights reserved.
#
# 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.
#
#
# Authors:
#       Chuk Orakwue <chuk.orakwue@huawei.com>

"""
    TaskState: Possible task states from linux/sched.h
    Task: Runnable thread (process) executed on CPU.
"""
from six import integer_types
from collections import namedtuple, defaultdict
from .common import ConstantBase

TaskBase = namedtuple("Task",
    [
    'name', # Task name
    'pid', # Process pid
    'prio', # Task priority
    'tgid', # Thread group ID
    'ppid', # Parent process ID
    ]
)

class TaskState(ConstantBase):
    RUNNING = ()
    SLEEPING = ()
    UNINTERRUPTIBLE = ()
    STOPPED = ()

    EXIT_DEAD = ()
    EXIT_ZOMBIE = ()

    TASK_DEAD = ()
    TASK_WAKEKILL = ()
    TASK_WAKING = ()
    TASK_PARKED = ()

    TRACED = ()
    UNKNOWN = ()

    # based on sched_wakeup
    RUNNABLE = ()

# linux/sched.h
# /trace/events/sched.h
TaskStateMapping = defaultdict(lambda: TaskState.UNKNOWN)
state_dict = {
    'S' : TaskState.SLEEPING,
    'D' : TaskState.UNINTERRUPTIBLE,
    'T' : TaskState.STOPPED,
    't' : TaskState.TRACED,
    'Z' : TaskState.EXIT_DEAD,
    'X' : TaskState.EXIT_ZOMBIE,
    'x' : TaskState.TASK_DEAD,
    'K' : TaskState.TASK_WAKEKILL,
    'W' : TaskState.TASK_WAKING,
    'P' : TaskState.TASK_PARKED,
    'R+': TaskState.RUNNABLE,
    'R' : TaskState.RUNNING,
}

TaskStateMapping.update(state_dict)

class Task(TaskBase):

    __slots__ = ()

    def __new__(cls, name, pid, prio=None, tgid=None, ppid=None, **kwargs):
        pid = int(pid) if pid else pid
        prio = int(prio) if prio else prio
        try:
            tgid = int(tgid)
        except:
            tgid = tgid

        return super(Task, cls).__new__(
            cls,
            name=name,
            pid=pid,
            prio=prio,
            tgid=tgid,
            ppid=ppid,
        )

    def __repr__(self):
        return "Task(name={}, pid={}, prio={})".format(
        self.name, self.pid, self.prio,
        )

    def __eq__(self, other):
        """
        Compare by PID only. Prority is subject to change (dynamically)
        as is task name (post-fork).
        """
        if isinstance(other, Task):
            return True if other.pid == self.pid else False
        elif isinstance(other, integer_types):
            return True if other == self.pid else False

        raise ValueError('{type} not supported'.format(type(other)))

    def __hash__(self):
        # IMPORTANT: Don't hash by name or priority
        # as those are subject to change in runtime
       return hash((self.pid))

    def affinity(self):
        """Return affinity if any, None otherwise"""
        try:
            return int(self.name.split('/')[-1][0])
        except:
            return

##normal scheduling policies
##    range: 0
##    (SCHED_OTHER, SCHED_IDLE, SCHED_BATCH)
##
##real-time policies
##    range: 1(low) to 99(high)
##    (SCHED_FIFO, SCHED_RR)
##
##PRIORITY = 20 + NICE
##NICE : -20 (high) to (20)
<commit_msg>Task: Oops in taskstate<commit_after>#!/usr/bin/python

# Copyright 2015 Huawei Devices USA Inc. All rights reserved.
#
# 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.
#
#
# Authors:
#       Chuk Orakwue <chuk.orakwue@huawei.com>

"""
    TaskState: Possible task states from linux/sched.h
    Task: Runnable thread (process) executed on CPU.
"""
from six import integer_types
from collections import namedtuple, defaultdict
from .common import ConstantBase

TaskBase = namedtuple("Task",
    [
    'name', # Task name
    'pid', # Process pid
    'prio', # Task priority
    'tgid', # Thread group ID
    'ppid', # Parent process ID
    ]
)

class TaskState(ConstantBase):
    RUNNING = ()
    SLEEPING = ()
    UNINTERRUPTIBLE = ()
    STOPPED = ()

    EXIT_DEAD = ()
    EXIT_ZOMBIE = ()

    TASK_DEAD = ()
    TASK_WAKEKILL = ()
    TASK_WAKING = ()
    TASK_PARKED = ()

    TRACED = ()
    UNKNOWN = ()

    # based on sched_wakeup
    RUNNABLE = ()

# linux/sched.h
# /trace/events/sched.h
TaskStateMapping = defaultdict(lambda: TaskState.UNKNOWN)
state_dict = {
    'S' : TaskState.SLEEPING,
    'D' : TaskState.UNINTERRUPTIBLE,
    'T' : TaskState.STOPPED,
    't' : TaskState.TRACED,
    'Z' : TaskState.EXIT_ZOMBIE,
    'X' : TaskState.EXIT_DEAD,
    'x' : TaskState.TASK_DEAD,
    'K' : TaskState.TASK_WAKEKILL,
    'W' : TaskState.TASK_WAKING,
    'P' : TaskState.TASK_PARKED,
    'R+': TaskState.RUNNABLE,
    'R' : TaskState.RUNNING,
}

TaskStateMapping.update(state_dict)

class Task(TaskBase):

    __slots__ = ()

    def __new__(cls, name, pid, prio=None, tgid=None, ppid=None, **kwargs):
        pid = int(pid) if pid else pid
        prio = int(prio) if prio else prio
        try:
            tgid = int(tgid)
        except:
            tgid = tgid

        return super(Task, cls).__new__(
            cls,
            name=name,
            pid=pid,
            prio=prio,
            tgid=tgid,
            ppid=ppid,
        )

    def __repr__(self):
        return "Task(name={}, pid={}, prio={})".format(
        self.name, self.pid, self.prio,
        )

    def __eq__(self, other):
        """
        Compare by PID only. Prority is subject to change (dynamically)
        as is task name (post-fork).
        """
        if isinstance(other, Task):
            return True if other.pid == self.pid else False
        elif isinstance(other, integer_types):
            return True if other == self.pid else False

        raise ValueError('{type} not supported'.format(type(other)))

    def __hash__(self):
        # IMPORTANT: Don't hash by name or priority
        # as those are subject to change in runtime
       return hash((self.pid))

    def affinity(self):
        """Return affinity if any, None otherwise"""
        try:
            return int(self.name.split('/')[-1][0])
        except:
            return

##normal scheduling policies
##    range: 0
##    (SCHED_OTHER, SCHED_IDLE, SCHED_BATCH)
##
##real-time policies
##    range: 1(low) to 99(high)
##    (SCHED_FIFO, SCHED_RR)
##
##PRIORITY = 20 + NICE
##NICE : -20 (high) to (20)
<|endoftext|>


<commit_before>###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    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/>.
#
###############################################################################
import os
import sys
import shutil
import parameters # Micronaet: configuration file
from datetime import datetime, timedelta

# -----------------------------------------------------------------------------
# Parameters:
# -----------------------------------------------------------------------------
demo = True #parameters.demo  # TODO change
input_folders = parameters.input_folders

print '''
Setup parameters:
    Demo: %s
    Input folders: %s
    ''' % (
        demo,
        input_folders,
        )

# Check elements:
error = [] # Error database
warning = [] # Warning database
info = [] #Info database
log = [] # Log database

# -----------------------------------------------------------------------------
#                           READ ALL INPUT FOLDERS:
# -----------------------------------------------------------------------------
import pdb; pdb.set_trace()
for (key, path, extension, walk) in input_folders:
    path = os.path.expanduser(path)
    for root, dirs, files in os.walk(path):
        for f in files:
            rename_file = False

            if f.startswith('.'):
                log.append('NO|Temp file not used|%s|' % f)
                continue
            if '.' not in f:
                log.append('NO|No extension (no .)|%s|' % f)
                continue

            # -----------------------------------------------------------------
            # Parse filename:
            # -----------------------------------------------------------------
            part = f.split('.')

            product = part[0].upper()
            prefix = ''
            version = ''
            if len(part) > 3:
                log.append('NO|File with extra dot|%s|' % f)
                continue
            elif len(part) == 3:
                version = part[1]

            # File name splitted:
            name = '.'.join(part[:-1]) # Take only first block at first dot!
            ext = part[-1]

            if not name:
                log.append('NO|No image file|%s|' % f)
                continue

            # -----------------------------------------------------------------
            # 1. Case problem in Extension (lower not jpeg:
            # -----------------------------------------------------------------
            if ext == ext.lower() and ext.lower() != 'jpeg':
                new_extension = ext # remain the same
            else:
                rename_file = True
                # Clean case error:
                new_extension = ext.lower()

                # Remove jpeg files:
                if new_extension == 'jpeg':
                    new_extension = 'jpg'

            # -----------------------------------------------------------------
            # 2. Case problem in name:
            # -----------------------------------------------------------------
            if name == name.upper():
                new_name = name
            else:
                rename_file = True
                new_name = name.upper()

            if new_name[:2] in ('ST', 'MT'):
                prefix = new_name[:2]
                new_name = new_name[2:] # remove prefix for analysis extra value
                product = product[2:]

            # -----------------------------------------------------------------
            # 3A. SX >> .001 change     or S.jpg >>> .jpg
            # -----------------------------------------------------------------
            if new_name[13:14] == 'S':
                rename_file = True
                extra_s = new_name[14:]
                if extra_s.isdigit():
                    new_name = '%s.%03d' % (
                        new_name[:13].strip('_'),
                        int(extra_s), # XXX < 3 char
                        )
                elif not extra_s:
                    new_name = new_name[:-1] # Remove S
                else:
                    log.append('NO|No SX format|%s|' % f)
                    continue

            # -----------------------------------------------------------------
            # 3B. SX >> .001 (13 char)
            # -----------------------------------------------------------------
            # Wrong SX format in position 13 instead of 14
            if new_name[12:13] == 'S':
                extra_s = new_name[13:]
                if extra_s.isdigit():
                    rename_file = True
                    new_name = '%s.%03d' % (
                        new_name[:12].strip('_'),
                        int(extra_s), # XXX < 3 char
                        )
                else:
                    pass # S(ingle) product

            # -----------------------------------------------------------------
            # 4. "NOME (1).jpg" >>> NOME.001.jpg
            # -----------------------------------------------------------------
            if '(' in new_name and ')' in new_name:
                extra_s = new_name[13:]
                extra_s = extra_s.replace('(', '').replace(')', '').strip()
                if extra_s.isdigit():
                    rename_file = True
                    new_name = '%s.%03d' % (
                        new_name[:13].strip('_'), # Remove extra _
                        int(extra_s), # XXX < 3 char
                        )

            # -----------------------------------------------------------------
            # 5. "NOME 1.jpg" >>> NOME.001.jpg
            # -----------------------------------------------------------------

            # -----------------------------------------------------------------
            # 6. + JUMPED
            # -----------------------------------------------------------------
            if '+' in new_name or '-' in new_name:
                log.append('NO|Jumped more product (+ or -)|%s|' % f)
                continue

            # -----------------------------------------------------------------
            # 7. + JUMPED
            # -----------------------------------------------------------------
            if 'COPIA' in new_name:
                log.append('NO|Jumped COPIA product|%s|' % f)
                continue

            # TODO case not managed: _COPIA, S1-0, NOME(1).jpg

            # -----------------------------------------------------------------
            # LAST. " " >>> "_"
            # -----------------------------------------------------------------
            if ' ' in new_name:
                rename_file = True
                new_name = new_name.replace(' ', '_')

            # -----------------------------------------------------------------
            # END: Rename procedure:
            # -----------------------------------------------------------------
            if not rename_file:
                if len(product) > 13:
                    log.append('NO|Check format (>13 char)|%s|' % f)
                    continue
                elif version and len(version) != 3:
                    log.append('NO|Check format (version 3 char)|%s|' % f)
                    continue
                elif version and not version.isdigit():
                    log.append('NO|Check format (version not #)|%s|' % f)
                    continue
                log.append('NO|No rename operation|%s|' % f)
                continue

            new_name = new_name.strip('_')
            new_name = '%s%s.%s' % (prefix, new_name, new_extension)
            from_file = os.path.join(root, f)
            to_file = os.path.join(root, new_name)
            if os.path.isfile(to_file):
                log.append('NO|Duplicated name|%s|%s' % (f, new_name))
                continue

            if demo:
                # Log operation:
                log.append('YES|RENAMED|%s|%s' % (f, new_name))
            else:
                # Rename operation:
                pass#shutil.move(from_file, to_file)
        break

f_log = open('./log.csv', 'w')
for item in sorted(log):
    f_log.write('%s\n' % item)
f_log.close()
# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>generic bugfix (fast message)<commit_after>###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    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/>.
#
###############################################################################
import os
import sys
import shutil
import parameters # Micronaet: configuration file
from datetime import datetime, timedelta

# -----------------------------------------------------------------------------
# Parameters:
# -----------------------------------------------------------------------------
demo = True #parameters.demo  # TODO change
input_folders = parameters.input_folders

print '''
Setup parameters:
    Demo: %s
    Input folders: %s
    ''' % (
        demo,
        input_folders,
        )

# Check elements:
error = [] # Error database
warning = [] # Warning database
info = [] #Info database
log = [] # Log database

# -----------------------------------------------------------------------------
#                           READ ALL INPUT FOLDERS:
# -----------------------------------------------------------------------------
import pdb; pdb.set_trace()
for (key, path, extension, walk) in input_folders:
    path = os.path.expanduser(path)
    for root, dirs, files in os.walk(path):
        for f in files:
            rename_file = False

            if f.startswith('.'):
                log.append('NO|Temp file not used|%s|%s' % (f, f))
                continue
            if '.' not in f:
                log.append('NO|No extension (no .)|%s|%s' % (f, f))
                continue

            # -----------------------------------------------------------------
            # Parse filename:
            # -----------------------------------------------------------------
            part = f.split('.')

            product = part[0].upper()
            prefix = ''
            version = ''
            if len(part) > 3:
                log.append('NO|File with extra dot|%s|%s' % (f, f))
                continue
            elif len(part) == 3:
                version = part[1]

            # File name splitted:
            name = '.'.join(part[:-1]) # Take only first block at first dot!
            ext = part[-1]

            if not name:
                log.append('NO|No image file|%s|%s' % (f, f))
                continue

            # -----------------------------------------------------------------
            # 1. Case problem in Extension (lower not jpeg:
            # -----------------------------------------------------------------
            if ext == ext.lower() and ext.lower() != 'jpeg':
                new_extension = ext # remain the same
            else:
                rename_file = True
                # Clean case error:
                new_extension = ext.lower()

                # Remove jpeg files:
                if new_extension == 'jpeg':
                    new_extension = 'jpg'

            # -----------------------------------------------------------------
            # 2. Case problem in name:
            # -----------------------------------------------------------------
            if name == name.upper():
                new_name = name
            else:
                rename_file = True
                new_name = name.upper()

            if new_name[:2] in ('ST', 'MT'):
                prefix = new_name[:2]
                new_name = new_name[2:] # remove prefix for analysis extra value
                product = product[2:]

            # -----------------------------------------------------------------
            # 3A. SX >> .001 change     or S.jpg >>> .jpg
            # -----------------------------------------------------------------
            if new_name[13:14] == 'S':
                rename_file = True
                extra_s = new_name[14:]
                if extra_s.isdigit():
                    new_name = '%s.%03d' % (
                        new_name[:13].strip('_'),
                        int(extra_s), # XXX < 3 char
                        )
                elif not extra_s:
                    new_name = new_name[:-1] # Remove S
                else:
                    log.append('NO|No SX format|%s|%s' % (f, f))
                    continue

            # -----------------------------------------------------------------
            # 3B. SX >> .001 (13 char)
            # -----------------------------------------------------------------
            # Wrong SX format in position 13 instead of 14
            if new_name[12:13] == 'S':
                extra_s = new_name[13:]
                if extra_s.isdigit():
                    rename_file = True
                    new_name = '%s.%03d' % (
                        new_name[:12].strip('_'),
                        int(extra_s), # XXX < 3 char
                        )
                else:
                    pass # S(ingle) product

            # -----------------------------------------------------------------
            # 4. "NOME (1).jpg" >>> NOME.001.jpg
            # -----------------------------------------------------------------
            if '(' in new_name and ')' in new_name:
                extra_s = new_name[13:]
                extra_s = extra_s.replace('(', '').replace(')', '').strip()
                if extra_s.isdigit():
                    rename_file = True
                    new_name = '%s.%03d' % (
                        new_name[:13].strip('_'), # Remove extra _
                        int(extra_s), # XXX < 3 char
                        )

            # -----------------------------------------------------------------
            # 5. "NOME 1.jpg" >>> NOME.001.jpg
            # -----------------------------------------------------------------

            # -----------------------------------------------------------------
            # 6. + JUMPED
            # -----------------------------------------------------------------
            if '+' in new_name or '-' in new_name:
                log.append('NO|Jumped more product (+ or -)|%s|%s' % (f, f))
                continue

            # -----------------------------------------------------------------
            # 7. + JUMPED
            # -----------------------------------------------------------------
            if 'COPIA' in new_name:
                log.append('NO|Jumped COPIA product|%s|%s' % (f, f))
                continue

            # -----------------------------------------------------------------
            # 8. + JUMPED
            # -----------------------------------------------------------------
            if  new_name[13:14] == 'A':
                log.append('NO|Environment image|%s|%s' % (f, f))
                continue

            # TODO case not managed: _COPIA, S1-0, NOME(1).jpg

            # -----------------------------------------------------------------
            # LAST. " " >>> "_"
            # -----------------------------------------------------------------
            if ' ' in new_name:
                rename_file = True
                new_name = new_name.replace(' ', '_')

            # -----------------------------------------------------------------
            # END: Rename procedure:
            # -----------------------------------------------------------------
            if not rename_file:
                if len(product) > 13:
                    log.append('NO|Check format (>13 char)|%s|%s' % (f, f))
                    continue
                elif version and len(version) != 3:
                    log.append(
                        'NO|Check format (version 3 char)|%s|%s' % (f, f))
                    continue
                elif version and not version.isdigit():
                    log.append(
                        'NO|Check format (version not #)|%s|%s' % (f, f))
                    continue
                log.append('NO|No rename operation|%s|%s' % (f, f))
                continue

            new_name = new_name.strip('_')
            new_name = '%s%s.%s' % (prefix, new_name, new_extension)
            from_file = os.path.join(root, f)
            to_file = os.path.join(root, new_name)
            if os.path.isfile(to_file):
                log.append('NO|Duplicated name|%s|%s' % (f, new_name))
                continue

            if demo:
                # Log operation:
                log.append('YES|RENAMED|%s|%s' % (f, new_name))
            else:
                # Rename operation:
                pass#shutil.move(from_file, to_file)
        break

f_log = open('./log.csv', 'w')
for item in sorted(log):
    f_log.write('%s\n' % item)
f_log.close()
# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before># Copyright 2020, The TensorFlow Authors.
#
# 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.

# Lint as: python3
"""Plotting functionality for membership inference attack analysis.

Functions to plot ROC curves and histograms as well as functionality to store
figures to colossus.
"""

from typing import Text, Iterable, Optional

import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics

def save_plot(figure: plt.Figure, path: Text, outformat='png'):
  """Store a figure to disk."""
  if path is not None:
    with open(path, 'wb') as f:
      figure.savefig(f, bbox_inches='tight', format=outformat)
    plt.close(figure)

def plot_curve_with_area(x: Iterable[float],
                         y: Iterable[float],
                         xlabel: Text = 'x',
                         ylabel: Text = 'y') -> plt.Figure:
  """Plot the curve defined by inputs and the area under the curve.

  All entries of x and y are required to lie between 0 and 1.
  For example, x could be recall and y precision, or x is fpr and y is tpr.

  Args:
    x: Values on x-axis (1d)
    y: Values on y-axis (must be same length as x)
    xlabel: Label for x axis
    ylabel: Label for y axis

  Returns:
    The matplotlib figure handle
  """
  fig = plt.figure()
  plt.plot([0, 1], [0, 1], 'k', lw=1.0)
  plt.plot(x, y, lw=2, label=f'AUC: {metrics.auc(x, y):.3f}')
  plt.xlabel(xlabel)
  plt.ylabel(ylabel)
  plt.legend()
  return fig

def plot_histograms(train: Iterable[float],
                    test: Iterable[float],
                    xlabel: Text = 'x',
                    thresh: Optional[float] = None) -> plt.Figure:
  """Plot histograms of training versus test metrics."""
  xmin = min(np.min(train), np.min(test))
  xmax = max(np.max(train), np.max(test))
  bins = np.linspace(xmin, xmax, 100)
  fig = plt.figure()
  plt.hist(test, bins=bins, density=True, alpha=0.5, label='test', log='y')
  plt.hist(train, bins=bins, density=True, alpha=0.5, label='train', log='y')
  if thresh is not None:
    plt.axvline(thresh, c='r', label=f'threshold = {thresh:.3f}')
  plt.xlabel(xlabel)
  plt.ylabel('normalized counts (density)')
  plt.legend()
  return fig

def plot_roc_curve(roc_curve) -> plt.Figure:
  """Plot the ROC curve and the area under the curve."""
  return plot_curve_with_area(
      roc_curve.fpr, roc_curve.tpr, xlabel='FPR', ylabel='TPR')
<commit_msg>For MIA plotting, allow customized plotting function and set equal x and y aspects.<commit_after># Copyright 2020, The TensorFlow Authors.
#
# 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.

# Lint as: python3
"""Plotting functionality for membership inference attack analysis.

Functions to plot ROC curves and histograms as well as functionality to store
figures to colossus.
"""

from typing import Text, Iterable, Optional

import matplotlib.pyplot as plt
import numpy as np
from sklearn import metrics

def save_plot(figure: plt.Figure, path: Text, outformat='png'):
  """Store a figure to disk."""
  if path is not None:
    with open(path, 'wb') as f:
      figure.savefig(f, bbox_inches='tight', format=outformat)
    plt.close(figure)

def plot_curve_with_area(x: Iterable[float],
                         y: Iterable[float],
                         xlabel: Text = 'x',
                         ylabel: Text = 'y') -> plt.Figure:
  """Plot the curve defined by inputs and the area under the curve.

  All entries of x and y are required to lie between 0 and 1.
  For example, x could be recall and y precision, or x is fpr and y is tpr.

  Args:
    x: Values on x-axis (1d)
    y: Values on y-axis (must be same length as x)
    xlabel: Label for x axis
    ylabel: Label for y axis

  Returns:
    The matplotlib figure handle
  """
  fig = plt.figure()
  plt.plot([0, 1], [0, 1], 'k', lw=1.0)
  plt.plot(x, y, lw=2, label=f'AUC: {metrics.auc(x, y):.3f}')
  plt.xlabel(xlabel)
  plt.ylabel(ylabel)
  plt.gca().set_aspect('equal', adjustable='box')
  plt.legend()
  return fig

def plot_histograms(train: Iterable[float],
                    test: Iterable[float],
                    xlabel: Text = 'x',
                    thresh: Optional[float] = None) -> plt.Figure:
  """Plot histograms of training versus test metrics."""
  xmin = min(np.min(train), np.min(test))
  xmax = max(np.max(train), np.max(test))
  bins = np.linspace(xmin, xmax, 100)
  fig = plt.figure()
  plt.hist(test, bins=bins, density=True, alpha=0.5, label='test', log='y')
  plt.hist(train, bins=bins, density=True, alpha=0.5, label='train', log='y')
  if thresh is not None:
    plt.axvline(thresh, c='r', label=f'threshold = {thresh:.3f}')
  plt.xlabel(xlabel)
  plt.ylabel('normalized counts (density)')
  plt.legend()
  return fig

def plot_roc_curve(roc_curve, plot_func=plot_curve_with_area) -> plt.Figure:
  """Plot the ROC curve and the area under the curve."""
  return plot_func(
      roc_curve.fpr, roc_curve.tpr, xlabel='FPR', ylabel='TPR')
<|endoftext|>