import datetime

import pandas as pd

import great_expectations as gx
import great_expectations.jupyter_ux
{%- if batch_request %}
from great_expectations.core.batch import BatchRequest
from great_expectations.checkpoint import SimpleCheckpoint
{% else %}
from great_expectations.core.expectation_configuration import ExpectationConfiguration
from great_expectations.data_context.types.resource_identifiers import ExpectationSuiteIdentifier
{% endif -%}
from great_expectations.exceptions import DataContextError

context = gx.get_context()
{% if batch_request %}
batch_request = {{ batch_request }}
{% endif %}

# Feel free to change the name of your suite here. Renaming this will not remove the other one.
expectation_suite_name = "{{ suite_name }}"
try:
    suite = context.get_expectation_suite(expectation_suite_name=expectation_suite_name)
    print(f'Loaded ExpectationSuite "{suite.expectation_suite_name}" containing {len(suite.expectations)} expectations.')
except DataContextError:
    suite = context.add_expectation_suite(expectation_suite_name=expectation_suite_name)
    print(f'Created ExpectationSuite "{suite.expectation_suite_name}".')

{% if batch_request %}
validator = context.get_validator(
    batch_request=BatchRequest(**batch_request),
    expectation_suite_name=expectation_suite_name
)
column_names = [f'"{column_name}"' for column_name in validator.columns()]
print(f"Columns: {', '.join(column_names)}.")
validator.head(n_rows=5, fetch_all=False)
{% endif %}
