Skip to content

Commit

Permalink
add download and extraction operators
Browse files Browse the repository at this point in the history
Signed-off-by: ilyashn <ilyashn@il.ibm.com>
  • Loading branch information
ilyashnil committed Jan 21, 2024
1 parent 6cb5661 commit 549e3ee
Showing 1 changed file with 40 additions and 0 deletions.
40 changes: 40 additions & 0 deletions src/unitxt/operator.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
import re
import zipfile
from abc import abstractmethod
from dataclasses import field
from typing import Any, Dict, Generator, List, Optional

import requests

from .artifact import Artifact
from .dataclass import NonPositionalField
from .stream import MultiStream, Stream
Expand All @@ -12,6 +15,43 @@ class Operator(Artifact):
pass


class DownloadError(Exception):
def __init__(
self,
message,
):
self.__super__(message)

Check warning on line 23 in src/unitxt/operator.py

View check run for this annotation

Codecov / codecov/patch

src/unitxt/operator.py#L23

Added line #L23 was not covered by tests


class UnexpectedHttpCodeError(Exception):
def __init__(self, http_code):
self.__super__(f"unexpected http code {http_code}")

Check warning on line 28 in src/unitxt/operator.py

View check run for this annotation

Codecov / codecov/patch

src/unitxt/operator.py#L28

Added line #L28 was not covered by tests


class DownloadOperator(Operator):
source: str
target: str

def __call__(self):
try:
response = requests.get(self.source, allow_redirects=True)
except Exception as e:
raise DownloadError(f"Unabled to download {self.source}") from e
if response.status_code != 200:
raise UnexpectedHttpCodeError(response.status_code)
with open(self.target, "wb") as f:
f.write(response.content)

Check warning on line 43 in src/unitxt/operator.py

View check run for this annotation

Codecov / codecov/patch

src/unitxt/operator.py#L36-L43

Added lines #L36 - L43 were not covered by tests


class ZipExtractorOperator(Operator):
zip_file: str
target_dir: str

def __call__(self):
with zipfile.ZipFile(self.zip_file) as zf:
zf.extractall(self.target_dir)

Check warning on line 52 in src/unitxt/operator.py

View check run for this annotation

Codecov / codecov/patch

src/unitxt/operator.py#L51-L52

Added lines #L51 - L52 were not covered by tests


class OperatorError(Exception):
def __init__(self, exception: Exception, operators: List[Operator]):
super().__init__(
Expand Down

0 comments on commit 549e3ee

Please sign in to comment.