neo4j-imdb-extractor/relations.js at master · frieck/neo4j-imdb-extractor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

let http = require('https');
let gunzip = require('zlib').createGunzip();
let readline = require('readline');
let fs = require('fs');

var dir = __dirname + '/data';
if (!fs.existsSync(dir)){
    fs.mkdirSync(dir);
}

let url = 'https://datasets.imdbws.com/title.principals.tsv.gz';
let csvFilepath = __dirname + '/data/title.principals.csv';
let csvHeaders = ':START_ID,role:string[],:END_ID,:TYPE \n';


let stream = fs.createWriteStream(csvFilepath);
stream.once('open', function () {
    http.get(url, response => {
        let len = parseInt(response.headers['content-length'], 10);
        let cur = 0;
        var total = len / 1048576; //1048576 - bytes in  1Megabyte

        response.on('data', function (chunk) {
            cur += chunk.length;
            process.stdout.cursorTo(0);
            process.stdout.write('Downloading ' + (100.0 * cur / len).toFixed(2) + '%    ' + (cur / 1048576).toFixed(2) + 'mb of Total size: ' + total.toFixed(2) + 'mb');
        });

        response.on('end', function () {
            console.log('\nDownloading complete');
        });

        let lineReader = readline.createInterface({
            input: response.pipe(gunzip)
        });

        let n = 0;
        lineReader.on('line', (line) => {

            if (n === 0) {
                stream.write(csvHeaders);
            } else {
                let values = line.split('\t');
                let movieID = values[0];
                let personID = values[2];
                let category = values[3];
                let char = values[5].replace('\\N', ' ');

                if(char !== ' ') {
                    let nchar = '"';
                    eval(char).forEach(c => {
                        nchar += c
                            .replace(/\[/g, '(')
                            .replace(/\]/g, ')')
                            .replace(/"/g, '""');

                    })
                    char = nchar + '"';
                }

                let rel = '';
                switch(category) {
                    case 'self':
                    case 'actor':
                    case 'actress':
                        rel = 'ACTED_IN';
                        break;

                    case 'director':
                        rel = 'DIRECTED';
                        break;

                    case 'composer':
                        rel = 'COMPOSED_FOR';
                        break;

                    case 'producer':
                        rel = 'PRODUCED';
                        break;

                    case 'editor':
                        rel = 'EDITED';
                        break;

                    case 'cinematographer':
                        rel = 'FILMED';
                        break;

                    case 'writer':
                        rel = 'WROTE';
                        break;
                }

                if(rel !== '') {
                    stream.write(`${personID},${char},${movieID},${rel} \n`);
                }

            }
            n++;
        });

        lineReader.on('close', () => {
            console.log(n + ' lines processed.');
            stream.end();
        });

    }).on('error', function (e) {
        console.log('Error: ' + e.message);
    });
});