import os

from .base import DataContext
from ..dataset.sparkdf_dataset import SparkDFDataset

from pyspark.sql import SparkSession


class SparkCSVDataContext(DataContext):
    """For now, functions like PandasCSVDataContext
    """

    def __init__(self, options, *args, **kwargs):
        super(SparkCSVDataContext, self).__init__(options, *args, **kwargs)
        self.spark = SparkSession.builder.getOrCreate()

    def connect(self, options):
        self.directory = options

    def list_datasets(self):
        return os.listdir(self.directory)

    def get_dataset(self, dataset_name, *args, caching=False, **kwargs):
        # TODO: deal with passing kwargs to spark.read and don't hard code options
        df = self.spark.read.option('header', True).csv(os.path.join(self.directory, dataset_name))
        return SparkDFDataset(df, caching=caching)

@@ -634,119 +634,106 @@
     
-# TODO: deal with passing kwargs to spark.read and don't hard code options%0A        df = self.spark.
+reader = self.spark.read%0A        for option in kwargs.items():%0A            reader = 
 read
+er
 .option(
 'hea
@@ -732,23 +732,36 @@
 ion(
-'header', True)
+*option)%0A        df = reader
 .csv

allow for dataset specific read options in spark csv context