<commit_before>import json
import logging
import os
from collections import OrderedDict
from io import BytesIO

from reportlab.lib.enums import TA_CENTER
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ListStyle, ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.platypus import (
    Image, PageBreak, Paragraph, SimpleDocTemplate, Spacer,
)

from django.conf import settings
from django.utils import timezone
from django.utils.html import conditional_escape

from callisto_core.utils import api

logger = logging.getLogger(__name__)

def report_as_pdf(report, data, recipient):
    return PDFFullReport(
        report=report,
        report_data=data,
    ).generate_pdf_report(
        recipient=recipient,
        report_id=report.id,
    )

class MatchReportContent(object):
    '''
        Class to structure contact information collected
        from match submission form for report
    '''

    # This constructor is called with keyword arguments populated by
    # encrypted data. Existing arguments should not be removed or renamed,
    # and new arguments must have default values.
    def __init__(
        self, identifier, perp_name, email, phone,
        contact_name=None, voicemail=None, notes=None,
    ):
        self.identifier = conditional_escape(identifier)
        self.perp_name = conditional_escape(perp_name)
        self.contact_name = conditional_escape(contact_name)
        self.email = conditional_escape(email)
        self.phone = conditional_escape(phone)
        self.voicemail = conditional_escape(voicemail)
        self.notes = conditional_escape(notes)

    @classmethod
    def from_form(cls, form):
        return cls(
            identifier=form.cleaned_data.get('identifier'),
            perp_name=form.cleaned_data.get('perp_name'),
            contact_name=form.instance.report.contact_name,
            email=form.instance.report.contact_email,
            phone=form.instance.report.contact_phone,
            voicemail=form.instance.report.contact_voicemail,
            notes=form.instance.report.contact_notes,
        )

class NumberedCanvas(canvas.Canvas):

    def __init__(self, *args, **kwargs):
        canvas.Canvas.__init__(self, *args, **kwargs)
        self._saved_page_states = []

    def showPage(self):
        self._saved_page_states.append(dict(self.__dict__))
        self._startPage()

    def save(self):
        '''add page info to each page (page x of y)'''
        num_pages = len(self._saved_page_states)
        for state in self._saved_page_states:
            self.__dict__.update(state)
            self.draw_page_number(num_pages)
            canvas.Canvas.showPage(self)
        canvas.Canvas.save(self)

    def draw_page_number(self, page_count):
        width, height = letter
        margin = 0.66 * 72
        self.setFillColor('gray')
        self.drawRightString(
            width - margin, margin, "Page %d of %d" %
            (self._pageNumber, page_count))

class PDFReport(object):

    unselected = u'\u2610'
    selected = u'\u2717'
    free_text = u'\u2756'
    no_bullet = ' '
    report_title = "Report"

    @property
    def headline_style(self):
        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48
        return headline_style

    @property
    def subtitle_style(self):
        styles = getSampleStyleSheet()
        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER
        return subtitle_style

    def __init__(self):
        self.styles = self.set_up_styles()
        self.notes_style = self.styles['Notes']
        self.answers_style = self.styles['Answers']
        self.answers_list_style = self.styles['AnswersList']
        self.body_style = self.styles['Normal']
        self.section_title_style = self.styles['Heading4']
        self.report_title_style = self.styles['Heading3']
        self.pdf_elements = []

    def set_up_styles(self, *args, **kwargs):
        '''
        Helper formatting methods assume the following styles have been defined:
            * Normal: format_question, get_metadata_page
            * Notes: get_metadata_page
            * Answers: format_answer, get_metadata_page
            * AnswersList: format_answer_list, get_metadata_page
            * Heading4: get_metadata_page, PDFFullReport.generate_pdf_report, PDFMatchReport.generate_match_report
            * Heading3: get_metadata_page, PDFMatchReport.generate_match_report
        If you override this method but don't define these styles, you will need to also override init and potentially
        the corresponding formatting methods.
        :return: a ReportLab stylesheet
        '''
        def get_font_location(filename):
            return settings.APPS_DIR.path('fonts')(filename)

        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48

        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER

        report_title_style = styles["Heading3"]
        report_title_style.fontSize = 28
        report_title_style.leading = 36
        report_title_style.spaceBefore = 0
        report_title_style.alignment = TA_CENTER

        section_title_style = styles["Heading4"]
        section_title_style.fontSize = 18
        section_title_style.leading = 22
        section_title_style.spaceBefore = 20

        body_style = styles["Normal"]
        body_style.fontSize = 14
        body_style.leading = 17
        body_style.leftIndent = 25

        styles.add(ParagraphStyle(
            name='Notes',
            parent=styles['Normal'],
            fontSize=10,
            leading=14,
            leftIndent=45,
            alias='notes',
        ))
        styles.add(ListStyle(
            name='AnswersList',
        ))
        styles.add(ParagraphStyle(
            name='Answers',
            parent=styles['Normal'],
            leftIndent=0,
            alias='answers',
        ))

        return styles

    # FORMATTING HELPERS

    def add_question(self, question):
        self.pdf_elements.append(
            Paragraph(question, self.body_style),
        )
        self.pdf_elements.append(Spacer(1, 4))

    def add_answer_list(self, answers):
        for answer in answers:
            self.pdf_elements.append(
                Paragraph(answer, self.notes_style),
            )
            self.pdf_elements.append(Spacer(1, 1))

    def render_question(self, question, answers):
        self.add_question(question)
        self.add_answer_list(answers)

    def render_questions(self, report_data):
        for item in report_data:
            question, answers = item.popitem()
            self.render_question(question, answers)

    def get_header_footer(self, recipient):
        def func(canvas, doc):
            width, height = letter
            margin = 0.66 * 72
            canvas.saveState()
            canvas.setFillColor('gray')
            canvas.drawString(margin, height - margin, "CONFIDENTIAL")
            canvas.drawRightString(
                width - margin, height - margin, str(timezone.now()))
            canvas.drawString(
                margin, margin,
                f"Intended for: Title IX Coordinator {recipient}",
            )
            canvas.restoreState()
        return func

    @staticmethod
    def get_user_identifier(user):
        try:
            return user.email or user.username
        except AttributeError:
            return 'Anonymous User'

class ReportPageMixin(object):

    def report_pages(self, reports: list):
        pages = []
        for report in reports:
            pages.extend(self.report_page(report))
            pages.append(PageBreak())
        return pages

    def report_page(self, report):
        return [
            Paragraph(
                "Report",
                self.report_title_style,
            ),
            Paragraph(
                "Report Metadata",
                self.section_title_style,
            ),
            Paragraph(
                f'''
                    Reported by: {self.get_user_identifier(report.owner)}

                    Submitted on: {report.submitted_to_school}

                    Record Created: {report.added.strftime("%Y-%m-%d %H:%M")}

                ''',
                self.body_style,
            ),
            Paragraph(
                "Contact Preferences",
                self.section_title_style,
            ),
            Paragraph(
                f'''
                    Name: {report.contact_name or "None provided"}

                    Phone: {report.contact_phone}

                    Voicemail preferences: {report.contact_voicemail or "None provided"}

                    Email: {report.contact_email}

                    Notes on preferred contact time of day, gender of admin, etc.:
                ''',
                self.body_style,
            ),
            Paragraph(
                report.contact_notes or "None provided",
                self.notes_style,
            )
        ]

class MatchPageMixin(object):

    def match_pages(self, matches: list):
        pages = []
        for match in matches:
            pages.extend(self.match_page(match))
            pages.append(PageBreak())
        return pages

    def match_page(self, match):
        return [
        ]

class PDFUserReviewReport(
    PDFReport,
    ReportPageMixin,
    MatchPageMixin,
):

    title = 'Submitted and Matched Reports'

    def cover_page(self):
        return [
            Image(
                os.path.join(settings.BASE_DIR, api.NotificationApi.logo_path),
                3 * inch,
                3 * inch,
            ),
            Spacer(1, 18),
            Paragraph("CONFIDENTIAL", self.headline_style),
            Spacer(1, 30),
            Spacer(1, 40),
            Paragraph(self.title, self.subtitle_style),
            Spacer(1, 40),
            PageBreak(),
        ]

    @classmethod
    def generate(cls, pdf_input_data: dict):
        # setup
        self = cls()
        reports = pdf_input_data.get('reports', [])
        matches = pdf_input_data.get('matches', [])
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(self.cover_page())
        self.pdf_elements.extend(self.report_pages(reports))
        self.pdf_elements.extend(self.match_pages(matches))

        # teardown
        doc.build(
            self.pdf_elements,
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFFullReport(
    PDFReport,
    ReportPageMixin,
):

    def __init__(self, report, report_data):
        super().__init__()
        self.report = report
        self.report_data = report_data

    def generate_pdf_report(self, report_id, recipient):
        # setup
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )
        self.pdf_elements.extend(self.report_page(self.report))
        self.pdf_elements.append(Paragraph("Report Questions", self.section_title_style))
        self.render_questions(self.report_data)

        # teardown
        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFMatchReport(PDFReport):

    report_title = "Match Report"

    def __init__(self, matches, identifier):
        super().__init__()
        self.matches = matches
        self.identifier = identifier

    def generate_match_report(self, report_id, recipient):
        '''
        Generates PDF report about a discovered match.

        Args:
          report_id (str): id used to uniquely identify this report to receiving authority

        Returns:
          bytes: a PDF with the submitted perp information & contact information of the reporters for this match

        '''
        # setup :: matches_with_reports
        sorted_matches = sorted(
            self.matches,
            key=lambda m: m.added.strftime("%Y-%m-%d %H:%M"),
            reverse=True,
        )
        matches_with_reports = [
            (
                match,
                MatchReportContent(**json.loads(match.get_match(self.identifier))),
            )
            for match in self.sorted_matches
        ]
        # setup :: names
        names = ', '.join(OrderedDict.fromkeys([report.perp_name.strip(
        ) for _, report in matches_with_reports if report.perp_name]))
        if len(names) < 1:
            names = 'None provided'
        # setup :: pdf
        buffer = BytesIO()
        doc = SimpleDocTemplate(
            buffer, pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )
        self.pdf_elements.append(
            Paragraph(
                api.NotificationApi.report_title,
                self.report_title_style,
            )
        )
        self.pdf_elements.append(
            Paragraph(
                "Perpetrator",
                self.section_title_style,
            )
        )
        self.pdf_elements.append(
            Paragraph(
                f'Name(s): {names}
Matching identifier: {self.identifier}',
                self.body_style,
            ),
        )
        # content fill :: reports
        for idx, (match_report, match_report_content) in enumerate(matches_with_reports):
            self.pdf_elements.append(
                Paragraph(
                    "Report " + str(idx + 1),
                    self.section_title_style,
                )
            )

            report = match_report.report
            user = report.owner

            if report.submitted_to_school:
                sent_report = report.sentfullreport_set.first()
                report_id = sent_report.get_report_id() if sent_report else "Not found"
                is_submitted = f'''
                    Yes

                    Submitted to school on: {report.submitted_to_school}

                    Submitted report ID: {report_id}
                '''
            else:
                is_submitted = "No"

            self.pdf_elements.append(
                Paragraph(
                    f'''
                        Perpetrator name given: {match_report_content.perp_name or "None provided"}

                        Reported by: {self.get_user_identifier(user)}

                        Submitted to matching on: {match_report.added.strftime("%Y-%m-%d %H:%M")}

                        Record created: {report.added.strftime("%Y-%m-%d %H:%M")}

                        Full record submitted? {is_submitted}

                        Name: {match_report_content.contact_name or "None provided"}

                        Phone: {match_report_content.phone}

                        Voicemail preferences: {match_report_content.voicemail or "None provided"}

                        Email: {match_report_content.email}

                        Notes on preferred contact time of day, gender of admin, etc.:

                    ''',
                    self.body_style,
                ),
            )
            self.pdf_elements.append(
                Paragraph(
                    match_report_content.notes or "None provided",
                    self.notes_style,
                ),
            )

        # teardown
        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = buffer.getvalue()
        buffer.close()
        return result
<commit_msg>wrong var name<commit_after>import json
import logging
import os
from collections import OrderedDict
from io import BytesIO

from reportlab.lib.enums import TA_CENTER
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import ListStyle, ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.platypus import (
    Image, PageBreak, Paragraph, SimpleDocTemplate, Spacer,
)

from django.conf import settings
from django.utils import timezone
from django.utils.html import conditional_escape

from callisto_core.utils import api

logger = logging.getLogger(__name__)

def report_as_pdf(report, data, recipient):
    return PDFFullReport(
        report=report,
        report_data=data,
    ).generate_pdf_report(
        recipient=recipient,
        report_id=report.id,
    )

class MatchReportContent(object):
    '''
        Class to structure contact information collected
        from match submission form for report
    '''

    # This constructor is called with keyword arguments populated by
    # encrypted data. Existing arguments should not be removed or renamed,
    # and new arguments must have default values.
    def __init__(
        self, identifier, perp_name, email, phone,
        contact_name=None, voicemail=None, notes=None,
    ):
        self.identifier = conditional_escape(identifier)
        self.perp_name = conditional_escape(perp_name)
        self.contact_name = conditional_escape(contact_name)
        self.email = conditional_escape(email)
        self.phone = conditional_escape(phone)
        self.voicemail = conditional_escape(voicemail)
        self.notes = conditional_escape(notes)

    @classmethod
    def from_form(cls, form):
        return cls(
            identifier=form.cleaned_data.get('identifier'),
            perp_name=form.cleaned_data.get('perp_name'),
            contact_name=form.instance.report.contact_name,
            email=form.instance.report.contact_email,
            phone=form.instance.report.contact_phone,
            voicemail=form.instance.report.contact_voicemail,
            notes=form.instance.report.contact_notes,
        )

class NumberedCanvas(canvas.Canvas):

    def __init__(self, *args, **kwargs):
        canvas.Canvas.__init__(self, *args, **kwargs)
        self._saved_page_states = []

    def showPage(self):
        self._saved_page_states.append(dict(self.__dict__))
        self._startPage()

    def save(self):
        '''add page info to each page (page x of y)'''
        num_pages = len(self._saved_page_states)
        for state in self._saved_page_states:
            self.__dict__.update(state)
            self.draw_page_number(num_pages)
            canvas.Canvas.showPage(self)
        canvas.Canvas.save(self)

    def draw_page_number(self, page_count):
        width, height = letter
        margin = 0.66 * 72
        self.setFillColor('gray')
        self.drawRightString(
            width - margin, margin, "Page %d of %d" %
            (self._pageNumber, page_count))

class PDFReport(object):

    unselected = u'\u2610'
    selected = u'\u2717'
    free_text = u'\u2756'
    no_bullet = ' '
    report_title = "Report"

    @property
    def headline_style(self):
        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48
        return headline_style

    @property
    def subtitle_style(self):
        styles = getSampleStyleSheet()
        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER
        return subtitle_style

    def __init__(self):
        self.styles = self.set_up_styles()
        self.notes_style = self.styles['Notes']
        self.answers_style = self.styles['Answers']
        self.answers_list_style = self.styles['AnswersList']
        self.body_style = self.styles['Normal']
        self.section_title_style = self.styles['Heading4']
        self.report_title_style = self.styles['Heading3']
        self.pdf_elements = []

    def set_up_styles(self, *args, **kwargs):
        '''
        Helper formatting methods assume the following styles have been defined:
            * Normal: format_question, get_metadata_page
            * Notes: get_metadata_page
            * Answers: format_answer, get_metadata_page
            * AnswersList: format_answer_list, get_metadata_page
            * Heading4: get_metadata_page, PDFFullReport.generate_pdf_report, PDFMatchReport.generate_match_report
            * Heading3: get_metadata_page, PDFMatchReport.generate_match_report
        If you override this method but don't define these styles, you will need to also override init and potentially
        the corresponding formatting methods.
        :return: a ReportLab stylesheet
        '''
        def get_font_location(filename):
            return settings.APPS_DIR.path('fonts')(filename)

        styles = getSampleStyleSheet()
        headline_style = styles["Heading1"]
        headline_style.alignment = TA_CENTER
        headline_style.fontSize = 48

        subtitle_style = styles["Heading2"]
        subtitle_style.fontSize = 24
        subtitle_style.leading = 26
        subtitle_style.alignment = TA_CENTER

        report_title_style = styles["Heading3"]
        report_title_style.fontSize = 28
        report_title_style.leading = 36
        report_title_style.spaceBefore = 0
        report_title_style.alignment = TA_CENTER

        section_title_style = styles["Heading4"]
        section_title_style.fontSize = 18
        section_title_style.leading = 22
        section_title_style.spaceBefore = 20

        body_style = styles["Normal"]
        body_style.fontSize = 14
        body_style.leading = 17
        body_style.leftIndent = 25

        styles.add(ParagraphStyle(
            name='Notes',
            parent=styles['Normal'],
            fontSize=10,
            leading=14,
            leftIndent=45,
            alias='notes',
        ))
        styles.add(ListStyle(
            name='AnswersList',
        ))
        styles.add(ParagraphStyle(
            name='Answers',
            parent=styles['Normal'],
            leftIndent=0,
            alias='answers',
        ))

        return styles

    # FORMATTING HELPERS

    def add_question(self, question):
        self.pdf_elements.append(
            Paragraph(question, self.body_style),
        )
        self.pdf_elements.append(Spacer(1, 4))

    def add_answer_list(self, answers):
        for answer in answers:
            self.pdf_elements.append(
                Paragraph(answer, self.notes_style),
            )
            self.pdf_elements.append(Spacer(1, 1))

    def render_question(self, question, answers):
        self.add_question(question)
        self.add_answer_list(answers)

    def render_questions(self, report_data):
        for item in report_data:
            question, answers = item.popitem()
            self.render_question(question, answers)

    def get_header_footer(self, recipient):
        def func(canvas, doc):
            width, height = letter
            margin = 0.66 * 72
            canvas.saveState()
            canvas.setFillColor('gray')
            canvas.drawString(margin, height - margin, "CONFIDENTIAL")
            canvas.drawRightString(
                width - margin, height - margin, str(timezone.now()))
            canvas.drawString(
                margin, margin,
                f"Intended for: Title IX Coordinator {recipient}",
            )
            canvas.restoreState()
        return func

    @staticmethod
    def get_user_identifier(user):
        try:
            return user.email or user.username
        except AttributeError:
            return 'Anonymous User'

class ReportPageMixin(object):

    def report_pages(self, reports: list):
        pages = []
        for report in reports:
            pages.extend(self.report_page(report))
            pages.append(PageBreak())
        return pages

    def report_page(self, report):
        return [
            Paragraph(
                "Report",
                self.report_title_style,
            ),
            Paragraph(
                "Report Metadata",
                self.section_title_style,
            ),
            Paragraph(
                f'''
                    Reported by: {self.get_user_identifier(report.owner)}

                    Submitted on: {report.submitted_to_school}

                    Record Created: {report.added.strftime("%Y-%m-%d %H:%M")}

                ''',
                self.body_style,
            ),
            Paragraph(
                "Contact Preferences",
                self.section_title_style,
            ),
            Paragraph(
                f'''
                    Name: {report.contact_name or "None provided"}

                    Phone: {report.contact_phone}

                    Voicemail preferences: {report.contact_voicemail or "None provided"}

                    Email: {report.contact_email}

                    Notes on preferred contact time of day, gender of admin, etc.:
                ''',
                self.body_style,
            ),
            Paragraph(
                report.contact_notes or "None provided",
                self.notes_style,
            )
        ]

class MatchPageMixin(object):

    def match_pages(self, matches: list):
        pages = []
        for match in matches:
            pages.extend(self.match_page(match))
            pages.append(PageBreak())
        return pages

    def match_page(self, match):
        return [
        ]

class PDFUserReviewReport(
    PDFReport,
    ReportPageMixin,
    MatchPageMixin,
):

    title = 'Submitted and Matched Reports'

    def cover_page(self):
        return [
            Image(
                os.path.join(settings.BASE_DIR, api.NotificationApi.logo_path),
                3 * inch,
                3 * inch,
            ),
            Spacer(1, 18),
            Paragraph("CONFIDENTIAL", self.headline_style),
            Spacer(1, 30),
            Spacer(1, 40),
            Paragraph(self.title, self.subtitle_style),
            Spacer(1, 40),
            PageBreak(),
        ]

    @classmethod
    def generate(cls, pdf_input_data: dict):
        # setup
        self = cls()
        reports = pdf_input_data.get('reports', [])
        matches = pdf_input_data.get('matches', [])
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(self.cover_page())
        self.pdf_elements.extend(self.report_pages(reports))
        self.pdf_elements.extend(self.match_pages(matches))

        # teardown
        doc.build(
            self.pdf_elements,
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFFullReport(
    PDFReport,
    ReportPageMixin,
):

    def __init__(self, report, report_data):
        super().__init__()
        self.report = report
        self.report_data = report_data

    def generate_pdf_report(self, report_id, recipient):
        # setup
        report_buffer = BytesIO()
        doc = SimpleDocTemplate(
            report_buffer,
            pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )
        self.pdf_elements.extend(self.report_page(self.report))
        self.pdf_elements.append(Paragraph("Report Questions", self.section_title_style))
        self.render_questions(self.report_data)

        # teardown
        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = report_buffer.getvalue()
        report_buffer.close()
        return result

class PDFMatchReport(PDFReport):

    report_title = "Match Report"

    def __init__(self, matches, identifier):
        super().__init__()
        self.matches = matches
        self.identifier = identifier

    def generate_match_report(self, report_id, recipient):
        '''
        Generates PDF report about a discovered match.

        Args:
          report_id (str): id used to uniquely identify this report to receiving authority

        Returns:
          bytes: a PDF with the submitted perp information & contact information of the reporters for this match

        '''
        # setup :: matches_with_reports
        sorted_matches = sorted(
            self.matches,
            key=lambda m: m.added.strftime("%Y-%m-%d %H:%M"),
            reverse=True,
        )
        matches_with_reports = [
            (
                match,
                MatchReportContent(**json.loads(match.get_match(self.identifier))),
            )
            for match in sorted_matches
        ]
        # setup :: names
        names = ', '.join(OrderedDict.fromkeys([report.perp_name.strip(
        ) for _, report in matches_with_reports if report.perp_name]))
        if len(names) < 1:
            names = 'None provided'
        # setup :: pdf
        buffer = BytesIO()
        doc = SimpleDocTemplate(
            buffer, pagesize=letter,
            rightMargin=72, leftMargin=72,
            topMargin=72, bottomMargin=72,
        )

        # content fill
        self.pdf_elements.extend(
            api.NotificationApi.get_cover_page(
                report_id=report_id,
                recipient=recipient,
            ),
        )
        self.pdf_elements.append(
            Paragraph(
                api.NotificationApi.report_title,
                self.report_title_style,
            )
        )
        self.pdf_elements.append(
            Paragraph(
                "Perpetrator",
                self.section_title_style,
            )
        )
        self.pdf_elements.append(
            Paragraph(
                f'Name(s): {names}
Matching identifier: {self.identifier}',
                self.body_style,
            ),
        )
        # content fill :: reports
        for idx, (match_report, match_report_content) in enumerate(matches_with_reports):
            self.pdf_elements.append(
                Paragraph(
                    "Report " + str(idx + 1),
                    self.section_title_style,
                )
            )

            report = match_report.report
            user = report.owner

            if report.submitted_to_school:
                sent_report = report.sentfullreport_set.first()
                report_id = sent_report.get_report_id() if sent_report else "Not found"
                is_submitted = f'''
                    Yes

                    Submitted to school on: {report.submitted_to_school}

                    Submitted report ID: {report_id}
                '''
            else:
                is_submitted = "No"

            self.pdf_elements.append(
                Paragraph(
                    f'''
                        Perpetrator name given: {match_report_content.perp_name or "None provided"}

                        Reported by: {self.get_user_identifier(user)}

                        Submitted to matching on: {match_report.added.strftime("%Y-%m-%d %H:%M")}

                        Record created: {report.added.strftime("%Y-%m-%d %H:%M")}

                        Full record submitted? {is_submitted}

                        Name: {match_report_content.contact_name or "None provided"}

                        Phone: {match_report_content.phone}

                        Voicemail preferences: {match_report_content.voicemail or "None provided"}

                        Email: {match_report_content.email}

                        Notes on preferred contact time of day, gender of admin, etc.:

                    ''',
                    self.body_style,
                ),
            )
            self.pdf_elements.append(
                Paragraph(
                    match_report_content.notes or "None provided",
                    self.notes_style,
                ),
            )

        # teardown
        doc.build(
            self.pdf_elements,
            onFirstPage=self.get_header_footer(recipient),
            onLaterPages=self.get_header_footer(recipient),
            canvasmaker=NumberedCanvas,
        )
        result = buffer.getvalue()
        buffer.close()
        return result
<|endoftext|>


<commit_before># (C) Datadog, Inc. 2010-present
# All rights reserved
# Licensed under Simplified BSD License (see LICENSE)

import collections

import mock

# project
import datadog_checks.btrfs

btrfs_check = datadog_checks.btrfs.BTRFS('btrfs', {}, [{}])

def mock_get_usage():

    return [
        (1, 9672065024, 9093722112),
        (34, 33554432, 16384),
        (36, 805306368, 544276480),
        (562949953421312, 184549376, 0),
    ]

# Just return a single device so the psutil portion of the check doesn't fail
# The real data to check against is in mock_get_usage.
def get_mock_devices():

    device_tuple = collections.namedtuple('device_tuple', 'device mountpoint fstype opts')

    return [device_tuple(device='/dev/disk1', mountpoint='/', fstype='btrfs', opts='local,multilabel')]

@mock.patch('datadog_checks.btrfs.btrfs.psutil.disk_partitions', return_value=get_mock_devices())
@mock.patch('datadog_checks.btrfs.btrfs.BTRFS.get_usage', return_value=mock_get_usage())
def test_check(mock_get_usage, mock_device_list, aggregator):
    """
    Testing Btrfs check.
    """
    with mock.patch.object(btrfs_check, 'get_unallocated_space', return_value=None):
        btrfs_check.check({})

    aggregator.assert_metric('system.disk.btrfs.unallocated', count=0)

    aggregator.reset()
    with mock.patch.object(btrfs_check, 'get_unallocated_space', return_value=0):
        btrfs_check.check({})

    aggregator.assert_metric('system.disk.btrfs.total', count=4)
    aggregator.assert_metric('system.disk.btrfs.used', count=4)
    aggregator.assert_metric('system.disk.btrfs.free', count=4)
    aggregator.assert_metric('system.disk.btrfs.usage', count=4)
    aggregator.assert_metric('system.disk.btrfs.unallocated', count=1)

    aggregator.assert_all_metrics_covered()

class MockStruct:
    def __init__(self, side_effects, size):
        self.counter = 0
        self.side_effects = side_effects
        self.size = size

    def unpack(self, data):
        ret = self.side_effects[self.counter]
        self.counter += 1
        return ret

    def unpack_from(self, data, _):
        return self.unpack(data)

    def pack_into(self, *v):
        pass

@mock.patch('datadog_checks.btrfs.btrfs.psutil.disk_partitions', return_value=get_mock_devices())
@mock.patch('datadog_checks.btrfs.btrfs.fcntl.ioctl')
def test_get_usage(mock_ioctl, mock_device_list):
    datadog_checks.btrfs.btrfs.TWO_LONGS_STRUCT = MockStruct([(0, 4), (0, 4)], 2)
    datadog_checks.btrfs.btrfs.THREE_LONGS_STRUCT = MockStruct(mock_get_usage(), 2)

    usage = btrfs_check.get_usage('/')

    assert usage == mock_get_usage()
<commit_msg>Use dd_run_check (#10655)<commit_after># (C) Datadog, Inc. 2010-present
# All rights reserved
# Licensed under Simplified BSD License (see LICENSE)

import collections

import mock

# project
import datadog_checks.btrfs

btrfs_check = datadog_checks.btrfs.BTRFS('btrfs', {}, [{}])

def mock_get_usage():

    return [
        (1, 9672065024, 9093722112),
        (34, 33554432, 16384),
        (36, 805306368, 544276480),
        (562949953421312, 184549376, 0),
    ]

# Just return a single device so the psutil portion of the check doesn't fail
# The real data to check against is in mock_get_usage.
def get_mock_devices():

    device_tuple = collections.namedtuple('device_tuple', 'device mountpoint fstype opts')

    return [device_tuple(device='/dev/disk1', mountpoint='/', fstype='btrfs', opts='local,multilabel')]

@mock.patch('datadog_checks.btrfs.btrfs.psutil.disk_partitions', return_value=get_mock_devices())
@mock.patch('datadog_checks.btrfs.btrfs.BTRFS.get_usage', return_value=mock_get_usage())
def test_check(mock_get_usage, mock_device_list, aggregator, dd_run_check):
    """
    Testing Btrfs check.
    """
    with mock.patch.object(btrfs_check, 'get_unallocated_space', return_value=None):
        dd_run_check(btrfs_check)

    aggregator.assert_metric('system.disk.btrfs.unallocated', count=0)

    aggregator.reset()
    with mock.patch.object(btrfs_check, 'get_unallocated_space', return_value=0):
        btrfs_check.check({})

    aggregator.assert_metric('system.disk.btrfs.total', count=4)
    aggregator.assert_metric('system.disk.btrfs.used', count=4)
    aggregator.assert_metric('system.disk.btrfs.free', count=4)
    aggregator.assert_metric('system.disk.btrfs.usage', count=4)
    aggregator.assert_metric('system.disk.btrfs.unallocated', count=1)

    aggregator.assert_all_metrics_covered()

class MockStruct:
    def __init__(self, side_effects, size):
        self.counter = 0
        self.side_effects = side_effects
        self.size = size

    def unpack(self, data):
        ret = self.side_effects[self.counter]
        self.counter += 1
        return ret

    def unpack_from(self, data, _):
        return self.unpack(data)

    def pack_into(self, *v):
        pass

@mock.patch('datadog_checks.btrfs.btrfs.psutil.disk_partitions', return_value=get_mock_devices())
@mock.patch('datadog_checks.btrfs.btrfs.fcntl.ioctl')
def test_get_usage(mock_ioctl, mock_device_list):
    datadog_checks.btrfs.btrfs.TWO_LONGS_STRUCT = MockStruct([(0, 4), (0, 4)], 2)
    datadog_checks.btrfs.btrfs.THREE_LONGS_STRUCT = MockStruct(mock_get_usage(), 2)

    usage = btrfs_check.get_usage('/')

    assert usage == mock_get_usage()
<|endoftext|>


<commit_before>from topaz.modules.ffi.abstract_memory import W_AbstractMemoryObject
from topaz.module import ClassDef
from topaz.coerce import Coerce

from rpython.rtyper.lltypesystem import rffi
from rpython.rtyper.lltypesystem import lltype

import sys

NULLPTR = lltype.nullptr(rffi.VOIDP.TO)

def coerce_pointer(space, w_pointer):
    if isinstance(w_pointer, W_PointerObject):
        return w_pointer.ptr
    else:
        raise space.error(space.w_TypeError,
                          "%s is not an FFI::Pointer." % w_pointer)

setattr(Coerce, 'ffi_pointer', staticmethod(coerce_pointer))

def coerce_address(space, w_addressable):
    if space.is_kind_of(w_addressable, space.w_bignum):
        return Coerce.int(space, w_addressable)
    elif space.is_kind_of(w_addressable, space.w_fixnum):
        return Coerce.int(space, w_addressable)
    elif space.is_kind_of(w_addressable,
                          space.getclassfor(W_PointerObject)):
        w_address = space.send(w_addressable, 'address')
        return coerce_address(space, w_address)
    else:
        errmsg = ("can't convert %s into FFI::Pointer" %
                  space.getclass(w_addressable).name)
        raise space.error(space.w_TypeError, errmsg)

setattr(Coerce, 'ffi_address', staticmethod(coerce_address))

class W_PointerObject(W_AbstractMemoryObject):
    classdef = ClassDef('Pointer', W_AbstractMemoryObject.classdef)

    def __init__(self, space, klass=None):
        W_AbstractMemoryObject.__init__(self, space, klass)
        self.sizeof_type = 0
        self.sizeof_memory = 0

    def __deepcopy__(self, memo):
        obj = super(W_AbstractMemoryObject, self).__deepcopy__(memo)
        obj.ptr = self.ptr
        obj.sizeof_type = self.sizeof_type
        obj.sizeof_memory = self.sizeof_memory
        return obj

    @classdef.singleton_method('allocate')
    def singleton_method_allocate(self, space):
        return W_PointerObject(space)

    @classdef.method('initialize')
    def method_initialize(self, space, args_w):
        if len(args_w) == 1:
            address = coerce_address(space, args_w[0])
            return self._initialize(space, address)
        elif len(args_w) == 2:
            sizeof_type = Coerce.int(space, args_w[0])
            address = coerce_address(space, args_w[1])
            return self._initialize(space, address, sizeof_type)

    def _initialize(self, space, address, sizeof_type=1):
        W_AbstractMemoryObject.__init__(self, space)
        self.ptr = rffi.cast(rffi.VOIDP, address)
        self.sizeof_type = sizeof_type
        self.sizeof_memory = sys.maxint

    @classdef.setup_class
    def setup_class(cls, space, w_cls):
        w_null = space.send(w_cls, 'new', [space.newint(0)])
        space.set_const(w_cls, 'NULL', w_null)
        space.send(w_cls, 'alias_method', [space.newsymbol('to_i'),
                                           space.newsymbol('address')])
        space.send(w_cls, 'alias_method', [space.newsymbol('[]'),
                                           space.newsymbol('+')])

    @classdef.method('free')
    def method_free(self, space):
        lltype.free(self.ptr, flavor='raw')

    @classdef.method('null?')
    def method_null_p(self, space):
        return space.newbool(self.ptr == NULLPTR)

    @classdef.method('address')
    def method_address(self, space):
        addr = rffi.cast(lltype.Unsigned, self.ptr)
        return space.newint_or_bigint_fromunsigned(addr)

    @classdef.method('size')
    def method_size(self, space):
        return space.newint_or_bigint(self.sizeof_memory)

    @classdef.method('==')
    def method_eq(self, space, w_other):
        if isinstance(w_other, W_PointerObject):
            return space.newbool(self.ptr == w_other.ptr)
        else:
            return space.newbool(False)

    @classdef.method('+', other='int')
    def method_plus(self, space, other):
        ptr = rffi.ptradd(rffi.cast(rffi.CCHARP, self.ptr), other)
        ptr_val = rffi.cast(lltype.Unsigned, ptr)
        w_ptr_val = space.newint_or_bigint_fromunsigned(ptr_val)
        w_res = space.send(space.getclass(self), 'new', [w_ptr_val])
        return w_res

    @classdef.method('slice', size='int')
    def method_address(self, space, w_offset, size):
        w_pointer = space.send(self, '+', [w_offset])
        assert isinstance(w_pointer, W_PointerObject)
        w_pointer.sizeof_memory = size
        return w_pointer

    @classdef.method('order', endianness='symbol')
    def method_order(self, space, endianness):
        return space.newint(0)

    @classdef.method('autorelease=', val='bool')
    def method_autorelease_eq(self, space, val):
        self.autorelease = val
        return space.newbool(val)

    @classdef.method('autorelease?')
    def method_autorelease_p(self, space):
        return space.newbool(self.autorelease)

    #@classdef.method('free')
    #def method_free(self, space):
    #    # TODO: Free stuff self is pointing at here
    #    return self

    @classdef.method('type_size')
    def method_type_size(self, space):
        return space.newint(self.sizeof_type)
<commit_msg>rename<commit_after>from topaz.modules.ffi.abstract_memory import W_AbstractMemoryObject
from topaz.module import ClassDef
from topaz.coerce import Coerce

from rpython.rtyper.lltypesystem import rffi
from rpython.rtyper.lltypesystem import lltype

import sys

NULLPTR = lltype.nullptr(rffi.VOIDP.TO)

def coerce_pointer(space, w_pointer):
    if isinstance(w_pointer, W_PointerObject):
        return w_pointer.ptr
    else:
        raise space.error(space.w_TypeError,
                          "%s is not an FFI::Pointer." % w_pointer)

setattr(Coerce, 'ffi_pointer', staticmethod(coerce_pointer))

def coerce_address(space, w_addressable):
    if space.is_kind_of(w_addressable, space.w_bignum):
        return Coerce.int(space, w_addressable)
    elif space.is_kind_of(w_addressable, space.w_fixnum):
        return Coerce.int(space, w_addressable)
    elif space.is_kind_of(w_addressable,
                          space.getclassfor(W_PointerObject)):
        w_address = space.send(w_addressable, 'address')
        return coerce_address(space, w_address)
    else:
        errmsg = ("can't convert %s into FFI::Pointer" %
                  space.getclass(w_addressable).name)
        raise space.error(space.w_TypeError, errmsg)

setattr(Coerce, 'ffi_address', staticmethod(coerce_address))

class W_PointerObject(W_AbstractMemoryObject):
    classdef = ClassDef('Pointer', W_AbstractMemoryObject.classdef)

    def __init__(self, space, klass=None):
        W_AbstractMemoryObject.__init__(self, space, klass)
        self.sizeof_type = 0
        self.sizeof_memory = 0

    def __deepcopy__(self, memo):
        obj = super(W_AbstractMemoryObject, self).__deepcopy__(memo)
        obj.ptr = self.ptr
        obj.sizeof_type = self.sizeof_type
        obj.sizeof_memory = self.sizeof_memory
        return obj

    @classdef.singleton_method('allocate')
    def singleton_method_allocate(self, space):
        return W_PointerObject(space)

    @classdef.method('initialize')
    def method_initialize(self, space, args_w):
        if len(args_w) == 1:
            address = coerce_address(space, args_w[0])
            return self._initialize(space, address)
        elif len(args_w) == 2:
            sizeof_type = Coerce.int(space, args_w[0])
            address = coerce_address(space, args_w[1])
            return self._initialize(space, address, sizeof_type)

    def _initialize(self, space, address, sizeof_type=1):
        W_AbstractMemoryObject.__init__(self, space)
        self.ptr = rffi.cast(rffi.VOIDP, address)
        self.sizeof_type = sizeof_type
        self.sizeof_memory = sys.maxint

    @classdef.setup_class
    def setup_class(cls, space, w_cls):
        w_null = space.send(w_cls, 'new', [space.newint(0)])
        space.set_const(w_cls, 'NULL', w_null)
        space.send(w_cls, 'alias_method', [space.newsymbol('to_i'),
                                           space.newsymbol('address')])
        space.send(w_cls, 'alias_method', [space.newsymbol('[]'),
                                           space.newsymbol('+')])

    @classdef.method('free')
    def method_free(self, space):
        lltype.free(self.ptr, flavor='raw')

    @classdef.method('null?')
    def method_null_p(self, space):
        return space.newbool(self.ptr == NULLPTR)

    @classdef.method('address')
    def method_address(self, space):
        addr = rffi.cast(lltype.Unsigned, self.ptr)
        return space.newint_or_bigint_fromunsigned(addr)

    @classdef.method('size')
    def method_size(self, space):
        return space.newint_or_bigint(self.sizeof_memory)

    @classdef.method('==')
    def method_eq(self, space, w_other):
        if isinstance(w_other, W_PointerObject):
            return space.newbool(self.ptr == w_other.ptr)
        else:
            return space.newbool(False)

    @classdef.method('+', other='int')
    def method_plus(self, space, other):
        ptr = rffi.ptradd(rffi.cast(rffi.CCHARP, self.ptr), other)
        ptr_val = rffi.cast(lltype.Unsigned, ptr)
        w_ptr_val = space.newint_or_bigint_fromunsigned(ptr_val)
        w_res = space.send(space.getclass(self), 'new', [w_ptr_val])
        return w_res

    @classdef.method('slice', size='int')
    def method_slice(self, space, w_offset, size):
        w_pointer = space.send(self, '+', [w_offset])
        assert isinstance(w_pointer, W_PointerObject)
        w_pointer.sizeof_memory = size
        return w_pointer

    @classdef.method('order', endianness='symbol')
    def method_order(self, space, endianness):
        return space.newint(0)

    @classdef.method('autorelease=', val='bool')
    def method_autorelease_eq(self, space, val):
        self.autorelease = val
        return space.newbool(val)

    @classdef.method('autorelease?')
    def method_autorelease_p(self, space):
        return space.newbool(self.autorelease)

    #@classdef.method('free')
    #def method_free(self, space):
    #    # TODO: Free stuff self is pointing at here
    #    return self

    @classdef.method('type_size')
    def method_type_size(self, space):
        return space.newint(self.sizeof_type)
<|endoftext|>


<commit_before># Copyright 2017 The TensorFlow Authors. 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.
# ==============================================================================
"""Timeseries head."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re

from tensorflow.contrib.framework.python.ops import variables
from tensorflow.contrib.layers.python.layers import optimizers

from tensorflow.contrib.timeseries.python.timeseries import feature_keys

from tensorflow.python.estimator import estimator_lib
from tensorflow.python.estimator.canned import head as head_lib
from tensorflow.python.estimator.export import export_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.util import nest

def time_series_regression_head(model,
                                state_manager,
                                optimizer,
                                input_statistics_generator=None):
  """Creates a `_Head` for time series regression.

  Args:
    model: A model for time series regression.
    state_manager: A state manager.
    optimizer: An optimizer.
    input_statistics_generator: A input statistics generator.

  Returns:
    An instance of `_Head` for time series regression.
  """
  return _TimeSeriesRegressionHead(model, state_manager, optimizer,
                                   input_statistics_generator)

class _TimeSeriesRegressionHead(head_lib._Head):  # pylint:disable=protected-access
  """See `time_series_regression_head`."""

  def __init__(self,
               model,
               state_manager,
               optimizer,
               input_statistics_generator=None,
               name=None):
    self.model = model
    self.state_manager = state_manager
    self.optimizer = optimizer
    self.input_statistics_generator = input_statistics_generator
    self._name = name

  def _train_ops(self, features):
    """Add training ops to the graph."""
    with variable_scope.variable_scope("model"):
      model_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.TRAIN)

    train_op = optimizers.optimize_loss(
        model_outputs.loss,
        global_step=variables.get_global_step(),
        optimizer=self.optimizer,
        # Learning rate is set in the Optimizer object
        learning_rate=None)
    return estimator_lib.EstimatorSpec(
        loss=model_outputs.loss,
        mode=estimator_lib.ModeKeys.TRAIN,
        train_op=train_op)

  # TODO(terrytangyuan): suffix summary and metrics keys by `"/" + name`
  @property
  def name(self):
    return self._name

  # TODO(terrytangyuan): unused for now. Need to decouple
  # `state_manager.define_loss` to satisfy the extendable return signature of
  # `_Head.create_loss`.
  def create_loss(self, features, mode, logits, labels):
    """See `_Head`."""
    return None

  # TODO(terrytangyuan): check label dimension
  @property
  def logits_dimension(self):
    return None

  def _evaluate_ops(self, features):
    """Add ops for evaluation (aka filtering) to the graph."""
    with variable_scope.variable_scope("model"):
      model_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.EVAL)
    metrics = {}
    # Just output in-sample predictions for the last chunk seen
    for prediction_key, prediction_value in model_outputs.predictions.items():
      metrics[prediction_key] = _identity_metric_single(prediction_key,
                                                        prediction_value)
    metrics[feature_keys.FilteringResults.TIMES] = _identity_metric_single(
        feature_keys.FilteringResults.TIMES, model_outputs.prediction_times)
    metrics[feature_keys.FilteringResults.STATE_TUPLE] = (
        _identity_metric_nested(feature_keys.FilteringResults.STATE_TUPLE,
                                model_outputs.end_state))
    return estimator_lib.EstimatorSpec(
        loss=model_outputs.loss,
        mode=estimator_lib.ModeKeys.EVAL,
        eval_metric_ops=metrics,
        predictions={})

  def _predict_ops(self, features):
    """Add ops for prediction to the graph."""
    with variable_scope.variable_scope("model"):
      prediction = self.model.predict(features=features)
    prediction[feature_keys.PredictionResults.TIMES] = features[
        feature_keys.PredictionFeatures.TIMES]
    return estimator_lib.EstimatorSpec(
        predictions=prediction, mode=estimator_lib.ModeKeys.PREDICT)

  def _serving_ops(self, features):
    """Add ops for serving to the graph."""
    with variable_scope.variable_scope("model"):
      prediction_outputs = self.model.predict(features=features)
    with variable_scope.variable_scope("model", reuse=True):
      filtering_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.EVAL)

    return estimator_lib.EstimatorSpec(
        mode=estimator_lib.ModeKeys.PREDICT,
        export_outputs={
            feature_keys.SavedModelLabels.PREDICT:
                export_lib.PredictOutput(prediction_outputs),
            feature_keys.SavedModelLabels.FILTER:
                export_lib.PredictOutput(
                    state_to_dictionary(filtering_outputs.end_state))
        },
        # Likely unused, but it is necessary to return `predictions` to satisfy
        # the Estimator's error checking.
        predictions={})

  def _convert_feature_to_tensor(self, name, value):
    """Casts features to the correct dtype based on their name."""
    if name in [
        feature_keys.TrainEvalFeatures.TIMES,
        feature_keys.PredictionFeatures.TIMES
    ]:
      return math_ops.cast(value, dtypes.int64)
    if name == feature_keys.TrainEvalFeatures.VALUES:
      return math_ops.cast(value, self.model.dtype)
    if name == feature_keys.PredictionFeatures.STATE_TUPLE:
      return value  # Correct dtypes are model-dependent
    return ops.convert_to_tensor(value)

  def _gather_state(self, features):
    """Returns `features` with state packed, indicates if packing was done."""
    prefixed_state_re = re.compile(r"^" + feature_keys.State.STATE_PREFIX +
                                   r"_(\d+)$")
    numbered_state = []
    for key, tensor in features.items():
      search_result = prefixed_state_re.search(key)
      if search_result:
        numbered_state.append((int(search_result.group(1)), key, tensor))
    if not numbered_state:
      return features, False
    features = features.copy()
    for _, key, _ in numbered_state:
      del features[key]
    numbered_state.sort(key=lambda number, *_: number)
    features[feature_keys.State.STATE_TUPLE] = nest.pack_sequence_as(
        structure=self.model.get_start_state(),
        flat_sequence=[tensor for _, _, tensor in numbered_state])
    return features, True

  def create_estimator_spec(self, features, mode, labels=None):
    """Performs basic error checking and returns an EstimatorSpec."""
    with ops.name_scope("head"):
      if labels:
        raise ValueError(
            "The model received a `labels` dictionary, which is "
            "not supported. Pass '{}' and '{}' as "
            "features.".format(feature_keys.TrainEvalFeatures.TIMES,
                               feature_keys.TrainEvalFeatures.VALUES))
      del labels
      features = {
          name: self._convert_feature_to_tensor(name=name, value=value)
          for name, value in features.items()
      }
      if self.input_statistics_generator is not None:
        input_statistics = self.input_statistics_generator.initialize_graph(
            features, update_statistics=(mode == estimator_lib.ModeKeys.TRAIN))
      else:
        input_statistics = None
      self.model.initialize_graph(input_statistics=input_statistics)

      # _gather_state requires the model to have its graph initialized (so it
      # has access to the structure of the model's state)
      features, passed_flat_state = self._gather_state(features)
      if (mode == estimator_lib.ModeKeys.TRAIN or
          mode == estimator_lib.ModeKeys.EVAL):
        _check_train_eval_features(features, self.model)
      elif mode == estimator_lib.ModeKeys.PREDICT:
        _check_predict_features(features)
      else:
        raise ValueError("Unknown mode '{}' passed to model_fn.".format(mode))

      self.state_manager.initialize_graph(
          model=self.model, input_statistics=input_statistics)

      if mode == estimator_lib.ModeKeys.TRAIN:
        return self._train_ops(features)
      elif mode == estimator_lib.ModeKeys.EVAL:
        return self._evaluate_ops(features)
      elif mode == estimator_lib.ModeKeys.PREDICT and not passed_flat_state:
        return self._predict_ops(features)
      elif mode == estimator_lib.ModeKeys.PREDICT and passed_flat_state:
        # The mode is PREDICT, but we're actually in export_savedmodel for
        # serving. We want to return two graphs: one for filtering (state + data
        # -> state) and one for predicting (state -> prediction).
        return self._serving_ops(features)

def _check_feature_shapes_compatible_with(features,
                                          compatible_with_name,
                                          compatible_with_value,
                                          ignore=None):
  """Checks all features are compatible with the given time-like feature."""
  if ignore is None:
    ignore = set()
  for name, value in features.items():
    if name in ignore:
      continue
    feature_shape = value.get_shape()
    if feature_shape.ndims is None:
      continue
    if feature_shape.ndims < 2:
      raise ValueError(
          ("Features must have shape (batch dimension, window size, ...) "
           "(got rank {} for feature '{}')").format(feature_shape.ndims, name))
    if not feature_shape[:2].is_compatible_with(
        compatible_with_value.get_shape()):
      raise ValueError(
          ("Features must have shape (batch dimension, window size, ...) "
           "where batch dimension and window size match the "
           "'{times_feature}' feature (got shape {feature_shape} for "
           "feature '{feature_name}' but shape {times_shape} for feature "
           "'{times_feature}')").format(
               times_feature=compatible_with_name,
               feature_shape=feature_shape,
               feature_name=name,
               times_shape=compatible_with_value.get_shape()))

def _check_predict_features(features):
  """Raises errors if features are not suitable for prediction."""
  if feature_keys.PredictionFeatures.TIMES not in features:
    raise ValueError("Expected a '{}' feature for prediction.".format(
        feature_keys.PredictionFeatures.TIMES))
  if feature_keys.PredictionFeatures.STATE_TUPLE not in features:
    raise ValueError("Expected a '{}' feature for prediction.".format(
        feature_keys.PredictionFeatures.STATE_TUPLE))
  times_feature = features[feature_keys.PredictionFeatures.TIMES]
  if not times_feature.get_shape().is_compatible_with([None, None]):
    raise ValueError(
        ("Expected shape (batch dimension, window size) for feature '{}' "
         "(got shape {})").format(feature_keys.PredictionFeatures.TIMES,
                                  times_feature.get_shape()))
  _check_feature_shapes_compatible_with(
      features=features,
      compatible_with_name=feature_keys.PredictionFeatures.TIMES,
      compatible_with_value=times_feature,
      ignore=set([
          feature_keys.PredictionFeatures.STATE_TUPLE  # Model-dependent shapes
      ]))

def _check_train_eval_features(features, model):
  """Raise errors if features are not suitable for training/evaluation."""
  if feature_keys.TrainEvalFeatures.TIMES not in features:
    raise ValueError("Expected a '{}' feature for training/evaluation.".format(
        feature_keys.TrainEvalFeatures.TIMES))
  if feature_keys.TrainEvalFeatures.VALUES not in features:
    raise ValueError("Expected a '{}' feature for training/evaluation.".format(
        feature_keys.TrainEvalFeatures.VALUES))
  times_feature = features[feature_keys.TrainEvalFeatures.TIMES]
  if not times_feature.get_shape().is_compatible_with([None, None]):
    raise ValueError(
        ("Expected shape (batch dimension, window size) for feature '{}' "
         "(got shape {})").format(feature_keys.TrainEvalFeatures.TIMES,
                                  times_feature.get_shape()))
  values_feature = features[feature_keys.TrainEvalFeatures.VALUES]
  if not values_feature.get_shape().is_compatible_with(
      [None, None, model.num_features]):
    raise ValueError(
        ("Expected shape (batch dimension, window size, {num_features}) "
         "for feature '{feature_name}', since the model was configured "
         "with num_features={num_features} (got shape {got_shape})").format(
             num_features=model.num_features,
             feature_name=feature_keys.TrainEvalFeatures.VALUES,
             got_shape=times_feature.get_shape()))
  _check_feature_shapes_compatible_with(
      features=features,
      compatible_with_name=feature_keys.TrainEvalFeatures.TIMES,
      compatible_with_value=times_feature,
      ignore=set([
          feature_keys.State.STATE_TUPLE  # Model-dependent shapes
      ]))

def _identity_metric_single(name, input_tensor):
  """A metric which takes on its last updated value.

  This keeps evaluation metrics in sync with one another, since update ops are
  run separately from their result Tensors. Simply returning (input_tensor,
  no_op) as a metric with a value but no update means that a metric will come
  from a different batch of data than metrics which cache values in a Variable
  (e.g. the default loss metric).

  Args:
    name: A name for the metric.
    input_tensor: Any Tensor.
  Returns:
    A tuple of (value, update_op).
  """
  metric_variable = variable_scope.variable(
      name="{}_identity_metric".format(name),
      initial_value=array_ops.zeros([], dtype=input_tensor.dtype),
      collections=[ops.GraphKeys.LOCAL_VARIABLES],
      validate_shape=False)
  update_op = state_ops.assign(
      metric_variable, input_tensor, validate_shape=False)
  # This shape will be correct once the first update runs (but may be
  # incomplete, so is not helpful for initializing the variable).
  metric_variable.set_shape(input_tensor.get_shape())
  return (metric_variable.value(), update_op)

def _identity_metric_nested(name, input_tensors):
  """Create identity metrics for a nested tuple of Tensors."""
  update_ops = []
  value_tensors = []
  for tensor_number, tensor in enumerate(nest.flatten(input_tensors)):
    value_tensor, update_op = _identity_metric_single(
        name="{}_{}".format(name, tensor_number), input_tensor=tensor)
    update_ops.append(update_op)
    value_tensors.append(value_tensor)
  return (nest.pack_sequence_as(input_tensors, value_tensors),
          control_flow_ops.group(*update_ops))

def state_to_dictionary(state_tuple):
  """Flatten model state into a dictionary with string keys."""
  flattened = {}
  for state_number, state_value in enumerate(nest.flatten(state_tuple)):
    prefixed_state_name = "{}_{:02d}".format(feature_keys.State.STATE_PREFIX,
                                             state_number)
    flattened[prefixed_state_name] = state_value
  return flattened
<commit_msg>variables.get_global_step() is deprecated,  use `training_util.get_global_step()` instead<commit_after># Copyright 2017 The TensorFlow Authors. 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.
# ==============================================================================
"""Timeseries head."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import re

from tensorflow.python.training import training_util
from tensorflow.contrib.layers.python.layers import optimizers

from tensorflow.contrib.timeseries.python.timeseries import feature_keys

from tensorflow.python.estimator import estimator_lib
from tensorflow.python.estimator.canned import head as head_lib
from tensorflow.python.estimator.export import export_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.util import nest

def time_series_regression_head(model,
                                state_manager,
                                optimizer,
                                input_statistics_generator=None):
  """Creates a `_Head` for time series regression.

  Args:
    model: A model for time series regression.
    state_manager: A state manager.
    optimizer: An optimizer.
    input_statistics_generator: A input statistics generator.

  Returns:
    An instance of `_Head` for time series regression.
  """
  return _TimeSeriesRegressionHead(model, state_manager, optimizer,
                                   input_statistics_generator)

class _TimeSeriesRegressionHead(head_lib._Head):  # pylint:disable=protected-access
  """See `time_series_regression_head`."""

  def __init__(self,
               model,
               state_manager,
               optimizer,
               input_statistics_generator=None,
               name=None):
    self.model = model
    self.state_manager = state_manager
    self.optimizer = optimizer
    self.input_statistics_generator = input_statistics_generator
    self._name = name

  def _train_ops(self, features):
    """Add training ops to the graph."""
    with variable_scope.variable_scope("model"):
      model_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.TRAIN)

    train_op = optimizers.optimize_loss(
        model_outputs.loss,
        global_step=training_util.get_global_step(),
        optimizer=self.optimizer,
        # Learning rate is set in the Optimizer object
        learning_rate=None)
    return estimator_lib.EstimatorSpec(
        loss=model_outputs.loss,
        mode=estimator_lib.ModeKeys.TRAIN,
        train_op=train_op)

  # TODO(terrytangyuan): suffix summary and metrics keys by `"/" + name`
  @property
  def name(self):
    return self._name

  # TODO(terrytangyuan): unused for now. Need to decouple
  # `state_manager.define_loss` to satisfy the extendable return signature of
  # `_Head.create_loss`.
  def create_loss(self, features, mode, logits, labels):
    """See `_Head`."""
    return None

  # TODO(terrytangyuan): check label dimension
  @property
  def logits_dimension(self):
    return None

  def _evaluate_ops(self, features):
    """Add ops for evaluation (aka filtering) to the graph."""
    with variable_scope.variable_scope("model"):
      model_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.EVAL)
    metrics = {}
    # Just output in-sample predictions for the last chunk seen
    for prediction_key, prediction_value in model_outputs.predictions.items():
      metrics[prediction_key] = _identity_metric_single(prediction_key,
                                                        prediction_value)
    metrics[feature_keys.FilteringResults.TIMES] = _identity_metric_single(
        feature_keys.FilteringResults.TIMES, model_outputs.prediction_times)
    metrics[feature_keys.FilteringResults.STATE_TUPLE] = (
        _identity_metric_nested(feature_keys.FilteringResults.STATE_TUPLE,
                                model_outputs.end_state))
    return estimator_lib.EstimatorSpec(
        loss=model_outputs.loss,
        mode=estimator_lib.ModeKeys.EVAL,
        eval_metric_ops=metrics,
        predictions={})

  def _predict_ops(self, features):
    """Add ops for prediction to the graph."""
    with variable_scope.variable_scope("model"):
      prediction = self.model.predict(features=features)
    prediction[feature_keys.PredictionResults.TIMES] = features[
        feature_keys.PredictionFeatures.TIMES]
    return estimator_lib.EstimatorSpec(
        predictions=prediction, mode=estimator_lib.ModeKeys.PREDICT)

  def _serving_ops(self, features):
    """Add ops for serving to the graph."""
    with variable_scope.variable_scope("model"):
      prediction_outputs = self.model.predict(features=features)
    with variable_scope.variable_scope("model", reuse=True):
      filtering_outputs = self.state_manager.define_loss(
          self.model, features, estimator_lib.ModeKeys.EVAL)

    return estimator_lib.EstimatorSpec(
        mode=estimator_lib.ModeKeys.PREDICT,
        export_outputs={
            feature_keys.SavedModelLabels.PREDICT:
                export_lib.PredictOutput(prediction_outputs),
            feature_keys.SavedModelLabels.FILTER:
                export_lib.PredictOutput(
                    state_to_dictionary(filtering_outputs.end_state))
        },
        # Likely unused, but it is necessary to return `predictions` to satisfy
        # the Estimator's error checking.
        predictions={})

  def _convert_feature_to_tensor(self, name, value):
    """Casts features to the correct dtype based on their name."""
    if name in [
        feature_keys.TrainEvalFeatures.TIMES,
        feature_keys.PredictionFeatures.TIMES
    ]:
      return math_ops.cast(value, dtypes.int64)
    if name == feature_keys.TrainEvalFeatures.VALUES:
      return math_ops.cast(value, self.model.dtype)
    if name == feature_keys.PredictionFeatures.STATE_TUPLE:
      return value  # Correct dtypes are model-dependent
    return ops.convert_to_tensor(value)

  def _gather_state(self, features):
    """Returns `features` with state packed, indicates if packing was done."""
    prefixed_state_re = re.compile(r"^" + feature_keys.State.STATE_PREFIX +
                                   r"_(\d+)$")
    numbered_state = []
    for key, tensor in features.items():
      search_result = prefixed_state_re.search(key)
      if search_result:
        numbered_state.append((int(search_result.group(1)), key, tensor))
    if not numbered_state:
      return features, False
    features = features.copy()
    for _, key, _ in numbered_state:
      del features[key]
    numbered_state.sort(key=lambda number, *_: number)
    features[feature_keys.State.STATE_TUPLE] = nest.pack_sequence_as(
        structure=self.model.get_start_state(),
        flat_sequence=[tensor for _, _, tensor in numbered_state])
    return features, True

  def create_estimator_spec(self, features, mode, labels=None):
    """Performs basic error checking and returns an EstimatorSpec."""
    with ops.name_scope("head"):
      if labels:
        raise ValueError(
            "The model received a `labels` dictionary, which is "
            "not supported. Pass '{}' and '{}' as "
            "features.".format(feature_keys.TrainEvalFeatures.TIMES,
                               feature_keys.TrainEvalFeatures.VALUES))
      del labels
      features = {
          name: self._convert_feature_to_tensor(name=name, value=value)
          for name, value in features.items()
      }
      if self.input_statistics_generator is not None:
        input_statistics = self.input_statistics_generator.initialize_graph(
            features, update_statistics=(mode == estimator_lib.ModeKeys.TRAIN))
      else:
        input_statistics = None
      self.model.initialize_graph(input_statistics=input_statistics)

      # _gather_state requires the model to have its graph initialized (so it
      # has access to the structure of the model's state)
      features, passed_flat_state = self._gather_state(features)
      if (mode == estimator_lib.ModeKeys.TRAIN or
          mode == estimator_lib.ModeKeys.EVAL):
        _check_train_eval_features(features, self.model)
      elif mode == estimator_lib.ModeKeys.PREDICT:
        _check_predict_features(features)
      else:
        raise ValueError("Unknown mode '{}' passed to model_fn.".format(mode))

      self.state_manager.initialize_graph(
          model=self.model, input_statistics=input_statistics)

      if mode == estimator_lib.ModeKeys.TRAIN:
        return self._train_ops(features)
      elif mode == estimator_lib.ModeKeys.EVAL:
        return self._evaluate_ops(features)
      elif mode == estimator_lib.ModeKeys.PREDICT and not passed_flat_state:
        return self._predict_ops(features)
      elif mode == estimator_lib.ModeKeys.PREDICT and passed_flat_state:
        # The mode is PREDICT, but we're actually in export_savedmodel for
        # serving. We want to return two graphs: one for filtering (state + data
        # -> state) and one for predicting (state -> prediction).
        return self._serving_ops(features)

def _check_feature_shapes_compatible_with(features,
                                          compatible_with_name,
                                          compatible_with_value,
                                          ignore=None):
  """Checks all features are compatible with the given time-like feature."""
  if ignore is None:
    ignore = set()
  for name, value in features.items():
    if name in ignore:
      continue
    feature_shape = value.get_shape()
    if feature_shape.ndims is None:
      continue
    if feature_shape.ndims < 2:
      raise ValueError(
          ("Features must have shape (batch dimension, window size, ...) "
           "(got rank {} for feature '{}')").format(feature_shape.ndims, name))
    if not feature_shape[:2].is_compatible_with(
        compatible_with_value.get_shape()):
      raise ValueError(
          ("Features must have shape (batch dimension, window size, ...) "
           "where batch dimension and window size match the "
           "'{times_feature}' feature (got shape {feature_shape} for "
           "feature '{feature_name}' but shape {times_shape} for feature "
           "'{times_feature}')").format(
               times_feature=compatible_with_name,
               feature_shape=feature_shape,
               feature_name=name,
               times_shape=compatible_with_value.get_shape()))

def _check_predict_features(features):
  """Raises errors if features are not suitable for prediction."""
  if feature_keys.PredictionFeatures.TIMES not in features:
    raise ValueError("Expected a '{}' feature for prediction.".format(
        feature_keys.PredictionFeatures.TIMES))
  if feature_keys.PredictionFeatures.STATE_TUPLE not in features:
    raise ValueError("Expected a '{}' feature for prediction.".format(
        feature_keys.PredictionFeatures.STATE_TUPLE))
  times_feature = features[feature_keys.PredictionFeatures.TIMES]
  if not times_feature.get_shape().is_compatible_with([None, None]):
    raise ValueError(
        ("Expected shape (batch dimension, window size) for feature '{}' "
         "(got shape {})").format(feature_keys.PredictionFeatures.TIMES,
                                  times_feature.get_shape()))
  _check_feature_shapes_compatible_with(
      features=features,
      compatible_with_name=feature_keys.PredictionFeatures.TIMES,
      compatible_with_value=times_feature,
      ignore=set([
          feature_keys.PredictionFeatures.STATE_TUPLE  # Model-dependent shapes
      ]))

def _check_train_eval_features(features, model):
  """Raise errors if features are not suitable for training/evaluation."""
  if feature_keys.TrainEvalFeatures.TIMES not in features:
    raise ValueError("Expected a '{}' feature for training/evaluation.".format(
        feature_keys.TrainEvalFeatures.TIMES))
  if feature_keys.TrainEvalFeatures.VALUES not in features:
    raise ValueError("Expected a '{}' feature for training/evaluation.".format(
        feature_keys.TrainEvalFeatures.VALUES))
  times_feature = features[feature_keys.TrainEvalFeatures.TIMES]
  if not times_feature.get_shape().is_compatible_with([None, None]):
    raise ValueError(
        ("Expected shape (batch dimension, window size) for feature '{}' "
         "(got shape {})").format(feature_keys.TrainEvalFeatures.TIMES,
                                  times_feature.get_shape()))
  values_feature = features[feature_keys.TrainEvalFeatures.VALUES]
  if not values_feature.get_shape().is_compatible_with(
      [None, None, model.num_features]):
    raise ValueError(
        ("Expected shape (batch dimension, window size, {num_features}) "
         "for feature '{feature_name}', since the model was configured "
         "with num_features={num_features} (got shape {got_shape})").format(
             num_features=model.num_features,
             feature_name=feature_keys.TrainEvalFeatures.VALUES,
             got_shape=times_feature.get_shape()))
  _check_feature_shapes_compatible_with(
      features=features,
      compatible_with_name=feature_keys.TrainEvalFeatures.TIMES,
      compatible_with_value=times_feature,
      ignore=set([
          feature_keys.State.STATE_TUPLE  # Model-dependent shapes
      ]))

def _identity_metric_single(name, input_tensor):
  """A metric which takes on its last updated value.

  This keeps evaluation metrics in sync with one another, since update ops are
  run separately from their result Tensors. Simply returning (input_tensor,
  no_op) as a metric with a value but no update means that a metric will come
  from a different batch of data than metrics which cache values in a Variable
  (e.g. the default loss metric).

  Args:
    name: A name for the metric.
    input_tensor: Any Tensor.
  Returns:
    A tuple of (value, update_op).
  """
  metric_variable = variable_scope.variable(
      name="{}_identity_metric".format(name),
      initial_value=array_ops.zeros([], dtype=input_tensor.dtype),
      collections=[ops.GraphKeys.LOCAL_VARIABLES],
      validate_shape=False)
  update_op = state_ops.assign(
      metric_variable, input_tensor, validate_shape=False)
  # This shape will be correct once the first update runs (but may be
  # incomplete, so is not helpful for initializing the variable).
  metric_variable.set_shape(input_tensor.get_shape())
  return (metric_variable.value(), update_op)

def _identity_metric_nested(name, input_tensors):
  """Create identity metrics for a nested tuple of Tensors."""
  update_ops = []
  value_tensors = []
  for tensor_number, tensor in enumerate(nest.flatten(input_tensors)):
    value_tensor, update_op = _identity_metric_single(
        name="{}_{}".format(name, tensor_number), input_tensor=tensor)
    update_ops.append(update_op)
    value_tensors.append(value_tensor)
  return (nest.pack_sequence_as(input_tensors, value_tensors),
          control_flow_ops.group(*update_ops))

def state_to_dictionary(state_tuple):
  """Flatten model state into a dictionary with string keys."""
  flattened = {}
  for state_number, state_value in enumerate(nest.flatten(state_tuple)):
    prefixed_state_name = "{}_{:02d}".format(feature_keys.State.STATE_PREFIX,
                                             state_number)
    flattened[prefixed_state_name] = state_value
  return flattened
<|endoftext|>