add types to parameters

Author: carlinmack

nsdb/nsdb.py

@@ -53,15 +53,15 @@ def createDumpsFile(listOfDumps: str, wiki: str, dump: str):
         subprocess.run(["./download.sh", "https://dumps.wikimedia.org/", wiki, dump])
 
 
-def countLines(file) -> int:
+def countLines(file: str) -> int:
     """Returns the number of lines in a file using wc from bash"""
     wordCount = subprocess.check_output(["wc", "-l", file]).decode("utf-8")
     lines = int(wordCount.split(" ")[0])
 
     return lines
 
 
-def downloadFirstDump(listOfDumps, archivesDir, dumpsDir) -> str:
+def downloadFirstDump(listOfDumps: str, archivesDir: str, dumpsDir: str) -> str:
     """Downloads the first dump in dumps.txt if it is not already present
     in the dumps directory"""
 
@@ -86,7 +86,7 @@ def downloadFirstDump(listOfDumps, archivesDir, dumpsDir) -> str:
     return fileName
 
 
-def extractFile(fileName: str, archivesDir, dumpsDir):
+def extractFile(fileName: str, archivesDir: str, dumpsDir: str):
     """Unzip if not already extracted, delete if extracted
 
     Execution takes 5-15 minutes as a guideline"""
@@ -99,7 +99,9 @@ def extractFile(fileName: str, archivesDir, dumpsDir):
     return fileName[:-3]
 
 
-def splitFile(fileName, queue, dumpsDir, partitionsDir, numPartitions):
+def splitFile(
+    fileName: str, queue, dumpsDir: str, partitionsDir: str, numPartitions: int
+):
     """Split a dump into a number of partitions"""
     database, cursor = Database.connect()
 
@@ -116,7 +118,7 @@ def splitFile(fileName, queue, dumpsDir, partitionsDir, numPartitions):
     database.close()
 
 
-def checkDiskSpace(dataDir):
+def checkDiskSpace(dataDir: str) -> int:
     """Returns the size of the data directory"""
     try:
         space = int(
@@ -188,7 +190,7 @@ def markLongRunningJobsAsError():
     database.close()
 
 
-def removeDoneJobs(partitionsDir):
+def removeDoneJobs(partitionsDir: str):
     """Remove partitions that are completed"""
     query = "SELECT file_name FROM partition WHERE status = 'done'"
     database, cursor = Database.connect()
@@ -242,7 +244,7 @@ def restartJobs():
     database.close()
 
 
-def main(parallelID=0, numParallel=1, dataDir="/bigtemp/ckm8gz/"):
+def main(parallelID: str = 0, numParallel: int = 1, dataDir: str = "/bigtemp/ckm8gz/"):
     """Download a list of dumps if it doesn't exist. If there are no dumps,
     download one and split it, then process the dump on multiple threads
 

nsdb/parse.py

@@ -11,7 +11,7 @@
 import time
 import traceback
 from datetime import datetime
-from typing import Tuple
+from typing import Tuple, List
 
 import mwreverts
 import mwxml
@@ -21,7 +21,7 @@
 import Database
 
 
-def multiprocess(partitionsDir, namespaces, queue, jobId):
+def multiprocess(partitionsDir: str, namespaces: List[int], queue, jobId: str):
     """Wrapper around process to call parse in a multiprocessing pool"""
     while True:
         i = queue.get()
@@ -35,7 +35,7 @@ def multiprocess(partitionsDir, namespaces, queue, jobId):
         time.sleep(10)
 
 
-def getDump(partitionsDir, cursor):
+def getDump(partitionsDir: str, cursor):
     """Returns the next dump to be parsed from the database
 
     Parameters
@@ -179,12 +179,12 @@ def parseNonTargetNamespace(
                 ),
             )
 
-    
     query = """UPDATE page (number_of_edits)
                 VALUES (%s) 
                 WHERE title=%s;"""
     cursor.execute(query, (pageEdits, title))
 
+
 def parseTargetNamespace(page, title: str, namespace: str, cursor, parallel: str):
     """Extracts features from each revision of a page into a database
 
@@ -431,7 +431,7 @@ def getDiff(old: str, new: str, parallel: str) -> Tuple[str, str]:
     return added, deleted
 
 
-def checkReverted(detector, revision, cursor, undidRevision, target):
+def checkReverted(detector, revision, cursor, undidRevision, target: bool):
     """Inserts reverted edits into the database for target namespace, otherwise 
     returns the user that was reverted"""
     reverted = detector.process(
@@ -575,7 +575,9 @@ def containsVulgarity(string: str) -> bool:
 
 
 def parse(
-    partitionsDir: str = "/bigtemp/ckm8gz/partitions/", namespaces=[1], parallel=""
+    partitionsDir: str = "/bigtemp/ckm8gz/partitions/",
+    namespaces: List[int] = [1],
+    parallel: str = "",
 ):
     """Selects the next dump from the database, extracts the features and
     imports them into several database tables.

nsdb/splitwiki.py

@@ -11,21 +11,21 @@
 from tqdm import trange
 
 
-def countLines(file) -> int:
+def countLines(file: str) -> int:
     """Returns the estimated number of lines in a dump using wcle.sh"""
     print("counting lines: ", file)
     lines = int(subprocess.check_output(["./wcle.sh", file]))
 
     return lines
 
 
-def addJobToDatabase(cursor, partitionName):
+def addJobToDatabase(cursor, partitionName: str):
     """Inserts partition into the database"""
     query = "INSERT INTO partition (file_name) VALUES (%s)"
     cursor.execute(query, (partitionName,))
 
 
-def addJobToQueue(queue, jobId):
+def addJobToQueue(queue, jobId: str):
     """Adds partition to the multiprocessing queue"""
     queue.put(jobId)
 
@@ -41,11 +41,11 @@ def addJobToQueue(queue, jobId):
 #     "-d", "--deletedump", default=False, is_flag=True,
 # )
 def split(
-    number=10,
-    inputFolder="/bigtemp/ckm8gz/dumps/",
-    outputFolder="/bigtemp/ckm8gz/partitions/",
-    deleteDump=True,
-    fileName="",
+    number: int = 10,
+    inputFolder: str = "/bigtemp/ckm8gz/dumps/",
+    outputFolder: str = "/bigtemp/ckm8gz/partitions/",
+    deleteDump: bool = True,
+    fileName: str = "",
     queue=0,
     cursor=0,
 ):
@@ -126,10 +126,10 @@ def split(
     i = 0
     inPage = False
     moreFile = True
-    
-    t = tqdm(total=number, desc=fileName, unit=" partition") # Initialise
+
+    t = tqdm(total=number, desc=fileName, unit=" partition")  # Initialise
     index = 0
-        
+
     with open(file) as inFile:
         while True:
             if not moreFile: