Overload DataFrame.drop: sequences must be *str (#377)

* Overload DataFrame.drop: sequences must be *str

Author: oliverw1

test-data/unit/sql-dataframe.test

@@ -22,3 +22,18 @@ df.sample(withReplacement=False) # E: No overload variant of "sample" of "DataFr
 [out]
 
 
+[case dropColumns]
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import lit, col
+spark = SparkSession.builder.getOrCreate()
+df = spark.range(1)
+df.drop("id")
+df.drop("id", "foo")
+df.drop(df.id)
+
+df.drop(col("id"), col("foo"))  # E: No overload variant of "drop" of "DataFrame" matches argument types "Column", "Column" \
+                                # N: Possible overload variant:     \
+                                # N:     def drop(self, *cols: str) -> DataFrame \
+                                # N:     <1 more non-matching overload not shown>
+
+[out]

third_party/3/pyspark/sql/dataframe.pyi

@@ -211,7 +211,10 @@ class DataFrame(PandasMapOpsMixin, PandasConversionMixin):
     ) -> DataFrame: ...
     def withColumn(self, colName: str, col: Column) -> DataFrame: ...
     def withColumnRenamed(self, existing: str, new: str) -> DataFrame: ...
-    def drop(self, *cols: ColumnOrName) -> DataFrame: ...
+    @overload
+    def drop(self, cols: ColumnOrName) -> DataFrame: ...
+    @overload
+    def drop(self, *cols: str) -> DataFrame: ...
     def toDF(self, *cols: ColumnOrName) -> DataFrame: ...
     def transform(self, func: Callable[[DataFrame], DataFrame]) -> DataFrame: ...
     @overload