text: from bz2 import open as bz2_open
  from logging import getLogger
  from multiprocessing import cpu_count
  from multiprocessing.pool import Pool
  from pathlib import Path
  from pickle import dump, load
  from tarfile import open as tarfile_open
  from tempfile import TemporaryDirectory
  from time import time
  from typing import Any, Dict, List
  
  from asdf import AsdfFile, open as asdf_open
  from requests import get as requests_get
  
  from mloncode.data.instance import Instance
  from mloncode.datasets.dataset import Dataset
  from mloncode.parsing.parser import Nodes, Parser, ParsingException
  
  
  class RepositoryDataset(Dataset):
  
      _logger = getLogger(__name__)
  
      def __init__(
          self,
          *,
          root_download_dir: str,
          root_parse_dir: str,
          root_tensor_dir: str,
          user_name: str,
          repo_name: str,
          version: str,
          instance: Instance,
          parser: Parser,
          bblfsh_endpoint: str = "0.0.0.0:9999",
          n_workers: int = cpu_count(),
          pickle_protocol: int = 4,
      ) -> None:
          # Output directories.
          self.root_download_dir = Path(root_download_dir).expanduser().resolve()
          self.root_parse_dir = Path(root_parse_dir).expanduser().resolve()
          self.root_tensor_dir = Path(root_tensor_dir).expanduser().resolve()
          # Downloading parameters.
          self.user_name = user_name
          self.repo_name = repo_name
          self.version = version
          # Parsing parameters.
          self.bblfsh_endpoint = bblfsh_endpoint
          # Tensorizing parameters.
          self.instance = instance
          self.parser = parser
          self.n_workers = n_workers
          self.pickle_protocol = pickle_protocol
          self.canonical_name = f"{self.user_name}-{self.repo_name}-{self.version}"
          self.download_path = self.root_download_dir / f"{self.canonical_name}.tar.gz"
          self.parse_dir = self.root_parse_dir / self.canonical_name
          self.tensor_dir = self.root_tensor_dir / self.canonical_name
  
      def download(self) -> None:
          self.root_download_dir.mkdir(parents=True, exist_ok=True)
          if self.download_path.is_file():
              return
          url = (
              f"https://github.com/{self.user_name}/{self.repo_name}/tarball/"
              f"{self.version}"
          )
          try:
              with requests_get(url, stream=True) as response:
                  response.raise_for_status()
                  with self.download_path.open("wb") as fh:
                      for chunk in response.iter_content(chunk_size=8192):
                          if chunk:
                              fh.write(chunk)
                  response.raise_for_status()
          except Exception as e:
              self.download_path.unlink()
              raise e
  
      def pre_tensorize(self) -> None:
          self._logger.info(f"Parsing {self.download_path.name}")
          self.parse_dir.mkdir(parents=True, exist_ok=True)
          if (self.parse_dir / "finished").is_file():
              return
          with TemporaryDirectory(prefix="mloncode-") as temporary_directory_name:
              with tarfile_open(self.download_path) as tar:
                  tar.extractall(path=temporary_directory_name)
              children_path = list(Path(temporary_directory_name).iterdir())
              assert len(children_path) == 1, "Expected a root directory only in the tar."
              repository_root_path = children_path[0].resolve(strict=True)
              for file_path in repository_root_path.rglob("*.js"):
                  file_path_relative = file_path.relative_to(repository_root_path)
                  try:
                      start = time()
                      self._logger.debug(f"Parsing {file_path_relative}")
                      nodes = self.parser.parse(repository_root_path, file_path_relative)
                      self._logger.debug(
                          f"Parsed  {file_path_relative} "
                          f"into {len(nodes.nodes)} nodes "
                          f"in {(time() - start) * 1000:.2f}ms"
                      )
                  except ParsingException:
                      continue
                  output_subdirectory = (
                      self.parse_dir / "asdf" / file_path_relative.parent
                  )
                  output_subdirectory.mkdir(parents=True, exist_ok=True)
                  with (output_subdirectory / file_path.with_suffix(".asdf").name).open(
                      "wb"
                  ) as fh:
                      af = AsdfFile(
                          dict(nodes=nodes.to_tree(file_path.read_text(encoding="utf-8")))
                      )
                      af.write_to(fh, all_array_compression="bzp2")
          (self.parse_dir / "finished").touch()
  
      def tensorize(self) -> None:
          try:
              # Cannot use multiprocessing if parser is an attribute of self.
              # Hence this hack with try, finally and a backup/restore of self.parser.
              parser = self.parser
              del self.parser
              finished_marker = self.tensor_dir / "finished"
              if not finished_marker.is_file():
                  self._logger.info(f"Pre-tensorizing {self.canonical_name}")
                  for file_path in self.parse_dir.rglob("*.asdf"):
                      with asdf_open(str(file_path)) as af:
                          nodes_instance = Nodes.from_tree(af.tree["nodes"])
                          self.instance.index(nodes_instance)
                  self._logger.info(f"Pre-tensorized  {self.canonical_name}")
                  self._logger.info(f"Tensorizing {self.canonical_name}")
                  with Pool(self.n_workers) as pool:
                      pool.map(
                          self._tensorize_worker,
                          [
                              p.relative_to(self.parse_dir)
                              for p in self.parse_dir.rglob("*.asdf")
                          ],
                      )
                  self._logger.info(f"Tensorized  {self.canonical_name}")
              self.pickles = list(
                  sorted((self.tensor_dir / "pickle").rglob("*.pickle.bz2"), key=str)
              )
              finished_marker.touch()
          finally:
              self.parser = parser
  
      def collate(self, tensors: List[Dict[str, Any]]) -> Dict[str, Any]:
          return self.instance.collate(tensors)
  
      def __getitem__(self, index: int) -> Dict[str, Any]:
          with bz2_open(self.pickles[index], "rb") as fh:
              return load(fh)
  
      def __len__(self) -> int:
          return len(self.pickles)
  
      def _tensorize_worker(self, file_path: Path) -> None:
          self._logger.debug(f"Tensorizing {file_path}")
          with asdf_open(str(self.parse_dir / file_path)) as af:
              nodes_instance = Nodes.from_tree(af.tree["nodes"])
          tensors = self.instance.tensorize(nodes_instance)
          output_dir = (self.tensor_dir / "pickle" / file_path).parent
          output_dir.mkdir(parents=True, exist_ok=True)
          with bz2_open(
              (output_dir / file_path.name).with_suffix(".pickle.bz2"), "wb"
          ) as fh:
              dump(tensors, fh, protocol=self.pickle_protocol)
          self._logger.debug(f"Tensorized  {file_path}")
  